Fix traffic stats not saved in direct boot mode

mobile/src/main/java/com/github/shadowsocks/App.kt

@@ -23,6 +23,7 @@ package com.github.shadowsocks
 import android.app.Application
 import android.app.NotificationChannel
 import android.app.NotificationManager
+import android.app.admin.DevicePolicyManager
 import android.content.BroadcastReceiver
 import android.content.Context
 import android.content.Intent
@@ -33,8 +34,8 @@ import android.content.res.Configuration
 import android.os.Build
 import android.os.Handler
 import android.os.Looper
+import android.os.UserManager
 import android.support.annotation.RequiresApi
-import android.support.v4.os.UserManagerCompat
 import android.support.v7.app.AppCompatDelegate
 import android.util.Log
 import android.widget.Toast
@@ -73,6 +74,10 @@ class App : Application() {
     private val tracker: Tracker by lazy { GoogleAnalytics.getInstance(deviceContext).newTracker(R.xml.tracker) }
     private val exceptionParser by lazy { StandardExceptionParser(this, null) }
     val info: PackageInfo by lazy { getPackageInfo(packageName) }
+    val directBootSupported by lazy {
+        Build.VERSION.SDK_INT >= 24 && getSystemService(DevicePolicyManager::class.java)
+            .storageEncryptionStatus == DevicePolicyManager.ENCRYPTION_STATUS_ACTIVE_PER_USER
+    }
 
     fun getPackageInfo(packageName: String) =
             packageManager.getPackageInfo(packageName, PackageManager.GET_SIGNATURES)!!
@@ -139,8 +144,9 @@ class App : Application() {
             app.track(e)
         }
 
-        // handle data restored
-        if (DataStore.directBootAware && UserManagerCompat.isUserUnlocked(this)) DirectBoot.update()
+        // handle data restored/crash
+        if (Build.VERSION.SDK_INT >= 24 && DataStore.directBootAware &&
+                (getSystemService(Context.USER_SERVICE) as UserManager).isUserUnlocked) DirectBoot.flushTrafficStats()
         TcpFastOpen.enabledAsync(DataStore.publicStore.getBoolean(Key.tfo, TcpFastOpen.sendEnabled))
         if (DataStore.publicStore.getLong(Key.assetUpdateTime, -1) != info.lastUpdateTime) {
             val assetManager = assets

mobile/src/main/java/com/github/shadowsocks/GlobalSettingsPreferenceFragment.kt

@@ -20,12 +20,12 @@
 
 package com.github.shadowsocks
 
-import android.app.admin.DevicePolicyManager
 import android.os.Build
 import android.os.Bundle
 import android.support.design.widget.Snackbar
 import android.support.v14.preference.SwitchPreference
 import android.support.v7.preference.Preference
+import com.github.shadowsocks.App.Companion.app
 import com.github.shadowsocks.bg.BaseService
 import com.github.shadowsocks.preference.DataStore
 import com.github.shadowsocks.utils.DirectBoot
@@ -39,20 +39,17 @@ class GlobalSettingsPreferenceFragment : PreferenceFragmentCompatDividers() {
         DataStore.initGlobal()
         addPreferencesFromResource(R.xml.pref_global)
         val boot = findPreference(Key.isAutoConnect) as SwitchPreference
-        val directBootAwareListener = Preference.OnPreferenceChangeListener { _, newValue ->
-            if (newValue as Boolean) DirectBoot.update() else DirectBoot.clean()
-            true
-        }
         boot.setOnPreferenceChangeListener { _, value ->
             BootReceiver.enabled = value as Boolean
-            directBootAwareListener.onPreferenceChange(null, DataStore.directBootAware)
+            true
         }
         boot.isChecked = BootReceiver.enabled
 
-        val dba = findPreference(Key.directBootAware)
-        if (Build.VERSION.SDK_INT >= 24 && requireContext().getSystemService(DevicePolicyManager::class.java)
-                .storageEncryptionStatus == DevicePolicyManager.ENCRYPTION_STATUS_ACTIVE_PER_USER)
-            dba.onPreferenceChangeListener = directBootAwareListener else dba.parent!!.removePreference(dba)
+        val canToggleLocked = findPreference(Key.directBootAware)
+        if (Build.VERSION.SDK_INT >= 24) canToggleLocked.setOnPreferenceChangeListener { _, newValue ->
+            if (app.directBootSupported && newValue as Boolean) DirectBoot.update() else DirectBoot.clean()
+            true
+        } else canToggleLocked.parent!!.removePreference(canToggleLocked)
 
         val tfo = findPreference(Key.tfo) as SwitchPreference
         tfo.isChecked = TcpFastOpen.sendEnabled

mobile/src/main/java/com/github/shadowsocks/bg/BaseService.kt

@@ -49,6 +49,7 @@ import okhttp3.OkHttpClient
 import okhttp3.Request
 import org.json.JSONObject
 import java.io.File
+import java.io.IOException
 import java.net.UnknownHostException
 import java.security.MessageDigest
 import java.util.*
@@ -145,22 +146,32 @@ object BaseService {
         }
 
         internal fun updateTrafficTotal(tx: Long, rx: Long) {
-            // this.profile may have host, etc. modified and thus a re-fetch is necessary (possible race condition)
-            val profile = ProfileManager.getProfile((profile ?: return).id) ?: return
-            profile.tx += tx
-            profile.rx += rx
-            ProfileManager.updateProfile(profile)
-            app.handler.post {
-                if (bandwidthListeners.isNotEmpty()) {
-                    val n = callbacks.beginBroadcast()
-                    for (i in 0 until n) {
-                        try {
-                            val item = callbacks.getBroadcastItem(i)
-                            if (bandwidthListeners.contains(item.asBinder())) item.trafficPersisted(profile.id)
-                        } catch (_: Exception) { }  // ignore
+            try {
+                // this.profile may have host, etc. modified and thus a re-fetch is necessary (possible race condition)
+                val profile = ProfileManager.getProfile((profile ?: return).id) ?: return
+                profile.tx += tx
+                profile.rx += rx
+                ProfileManager.updateProfile(profile)
+                app.handler.post {
+                    if (bandwidthListeners.isNotEmpty()) {
+                        val n = callbacks.beginBroadcast()
+                        for (i in 0 until n) {
+                            try {
+                                val item = callbacks.getBroadcastItem(i)
+                                if (bandwidthListeners.contains(item.asBinder())) item.trafficPersisted(profile.id)
+                            } catch (_: Exception) { }  // ignore
+                        }
+                        callbacks.finishBroadcast()
                     }
-                    callbacks.finishBroadcast()
                 }
+            } catch (e: IOException) {
+                if (!DataStore.directBootAware) throw e // we should only reach here because we're in direct boot
+                val profile = DirectBoot.getDeviceProfile()!!
+                profile.tx += tx
+                profile.rx += rx
+                profile.dirty = true
+                DirectBoot.update(profile)
+                DirectBoot.listenForUnlock()
             }
         }
 

mobile/src/main/java/com/github/shadowsocks/bg/TileService.kt

@@ -82,7 +82,7 @@ class TileService : BaseTileService(), ShadowsocksConnection.Interface {
     }
 
     override fun onClick() {
-        if (isLocked && !DataStore.directBootAware) unlockAndRun(this::toggle) else toggle()
+        if (isLocked && !DataStore.canToggleLocked) unlockAndRun(this::toggle) else toggle()
     }
 
     private fun toggle() {

mobile/src/main/java/com/github/shadowsocks/database/Profile.kt

@@ -147,6 +147,9 @@ class Profile : Serializable {
     @DatabaseField
     var plugin: String? = null
 
+    // not persisted in db, only used by direct boot
+    var dirty: Boolean = false
+
     val formattedAddress get() = (if (host.contains(":")) "[%s]:%d" else "%s:%d").format(host, remotePort)
     val formattedName get() = if (name.isNullOrEmpty()) formattedAddress else name!!
 

mobile/src/main/java/com/github/shadowsocks/database/ProfileManager.kt

@@ -20,11 +20,13 @@
 
 package com.github.shadowsocks.database
 
+import android.database.sqlite.SQLiteCantOpenDatabaseException
 import android.util.Log
 import com.github.shadowsocks.App.Companion.app
 import com.github.shadowsocks.ProfilesFragment
 import com.github.shadowsocks.preference.DataStore
 import com.github.shadowsocks.utils.DirectBoot
+import java.io.IOException
 import java.sql.SQLException
 
 /**
@@ -62,9 +64,11 @@ object ProfileManager {
     @Throws(SQLException::class)
     fun updateProfile(profile: Profile) = PrivateDatabase.profileDao.update(profile)
 
+    @Throws(IOException::class)
     fun getProfile(id: Int): Profile? = try {
         PrivateDatabase.profileDao.queryForId(id)
     } catch (ex: SQLException) {
+        if (ex.cause is SQLiteCantOpenDatabaseException) throw IOException(ex)
         Log.e(TAG, "getProfile", ex)
         app.track(ex)
         null
@@ -77,17 +81,21 @@ object ProfileManager {
         if (id == DataStore.profileId && DataStore.directBootAware) DirectBoot.clean()
     }
 
+    @Throws(IOException::class)
     fun getFirstProfile(): Profile? = try {
         PrivateDatabase.profileDao.query(PrivateDatabase.profileDao.queryBuilder().limit(1L).prepare()).singleOrNull()
     } catch (ex: SQLException) {
+        if (ex.cause is SQLiteCantOpenDatabaseException) throw IOException(ex)
         Log.e(TAG, "getFirstProfile", ex)
         app.track(ex)
         null
     }
 
+    @Throws(IOException::class)
     fun getAllProfiles(): List<Profile>? = try {
         PrivateDatabase.profileDao.query(PrivateDatabase.profileDao.queryBuilder().orderBy("userOrder", true).prepare())
     } catch (ex: SQLException) {
+        if (ex.cause is SQLiteCantOpenDatabaseException) throw IOException(ex)
         Log.e(TAG, "getAllProfiles", ex)
         app.track(ex)
         null

mobile/src/main/java/com/github/shadowsocks/preference/DataStore.kt

@@ -21,6 +21,7 @@
 package com.github.shadowsocks.preference
 
 import android.os.Binder
+import com.github.shadowsocks.App.Companion.app
 import com.github.shadowsocks.BootReceiver
 import com.github.shadowsocks.database.PrivateDatabase
 import com.github.shadowsocks.database.PublicDatabase
@@ -49,10 +50,8 @@ object DataStore {
             publicStore.putInt(Key.id, value)
             if (DataStore.directBootAware) DirectBoot.update()
         }
-    /**
-     * Setter is defined in MainActivity.onPreferenceDataStoreChanged.
-     */
-    val directBootAware: Boolean get() = BootReceiver.enabled && publicStore.getBoolean(Key.directBootAware) == true
+    val canToggleLocked: Boolean get() = publicStore.getBoolean(Key.directBootAware) == true
+    val directBootAware: Boolean get() = app.directBootSupported && canToggleLocked
     var serviceMode: String
         get() = publicStore.getString(Key.serviceMode) ?: Key.modeVpn
         set(value) = publicStore.putString(Key.serviceMode, value)

mobile/src/main/java/com/github/shadowsocks/utils/DirectBoot.kt

@@ -1,6 +1,10 @@
 package com.github.shadowsocks.utils
 
 import android.annotation.TargetApi
+import android.content.BroadcastReceiver
+import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
 import com.github.shadowsocks.App.Companion.app
 import com.github.shadowsocks.bg.BaseService
 import com.github.shadowsocks.database.Profile
@@ -12,8 +16,9 @@ import java.io.ObjectInputStream
 import java.io.ObjectOutputStream
 
 @TargetApi(24)
-object DirectBoot {
+object DirectBoot : BroadcastReceiver() {
     private val file = File(app.deviceContext.noBackupFilesDir, "directBootProfile")
+    private var registered = false
 
     fun getDeviceProfile(): Profile? = try {
         ObjectInputStream(file.inputStream()).use { it.readObject() as Profile }
@@ -24,8 +29,26 @@ object DirectBoot {
         File(app.deviceContext.noBackupFilesDir, BaseService.CONFIG_FILE).delete()
     }
 
-    fun update() {
-        val profile = ProfileManager.getProfile(DataStore.profileId)    // app.currentProfile will call this
-        if (profile == null) clean() else ObjectOutputStream(file.outputStream()).use { it.writeObject(profile) }
+    /**
+     * app.currentProfile will call this.
+     */
+    fun update(profile: Profile? = ProfileManager.getProfile(DataStore.profileId)) =
+            if (profile == null) clean() else ObjectOutputStream(file.outputStream()).use { it.writeObject(profile) }
+
+    fun flushTrafficStats() {
+        val profile = getDeviceProfile()
+        if (profile?.dirty == true) ProfileManager.updateProfile(profile)
+        update()
+    }
+
+    fun listenForUnlock() {
+        if (registered) return
+        app.registerReceiver(this, IntentFilter(Intent.ACTION_BOOT_COMPLETED))
+        registered = true
+    }
+    override fun onReceive(context: Context, intent: Intent) {
+        flushTrafficStats()
+        app.unregisterReceiver(this)
+        registered = false
     }
 }