[DataStore] Support backup data with compression

Bug: 325144964
Test: Manual tests
(cherry picked from https://googleplex-android-review.googlesource.com/q/commit:4303555ef342eea5dc3d50eb8f5bfe68639f9c91)
Merged-In: Ida5d71f2b39aeb271f4ad10e200f28dbfe5ed9c8
Change-Id: Ida5d71f2b39aeb271f4ad10e200f28dbfe5ed9c8
diff --git a/packages/SettingsLib/DataStore/Android.bp b/packages/SettingsLib/DataStore/Android.bp
index 1806544..9fafcab 100644
--- a/packages/SettingsLib/DataStore/Android.bp
+++ b/packages/SettingsLib/DataStore/Android.bp
@@ -14,4 +14,5 @@
         "androidx.core_core-ktx",
         "guava",
     ],
+    kotlincflags: ["-Xjvm-default=all"],
 }
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupCodec.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupCodec.kt
new file mode 100644
index 0000000..550645f
--- /dev/null
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupCodec.kt
@@ -0,0 +1,98 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.datastore
+
+import androidx.annotation.IntDef
+import java.io.InputStream
+import java.io.OutputStream
+import java.util.zip.Deflater
+import java.util.zip.DeflaterOutputStream
+import java.util.zip.InflaterInputStream
+
+/** Unique id of the codec. */
+@Target(AnnotationTarget.TYPE)
+@IntDef(
+    BackupCodecId.NO_OP.toInt(),
+    BackupCodecId.ZIP.toInt(),
+)
+@Retention(AnnotationRetention.SOURCE)
+annotation class BackupCodecId {
+    companion object {
+        /** Unknown reason of the change. */
+        const val NO_OP: Byte = 0
+        /** Data is updated. */
+        const val ZIP: Byte = 1
+    }
+}
+
+/** How to encode/decode the backup data. */
+interface BackupCodec {
+    /** Unique id of the codec. */
+    val id: @BackupCodecId Byte
+
+    /** Name of the codec. */
+    val name: String
+
+    /** Encodes the backup data. */
+    fun encode(outputStream: OutputStream): OutputStream
+
+    /** Decodes the backup data. */
+    fun decode(inputStream: InputStream): InputStream
+
+    companion object {
+        @JvmStatic
+        fun fromId(id: @BackupCodecId Byte): BackupCodec =
+            when (id) {
+                BackupCodecId.NO_OP -> BackupNoOpCodec()
+                BackupCodecId.ZIP -> BackupZipCodec.BEST_COMPRESSION
+                else -> throw IllegalArgumentException("Unknown codec id $id")
+            }
+    }
+}
+
+/** Codec without any additional encoding/decoding. */
+class BackupNoOpCodec : BackupCodec {
+    override val id
+        get() = BackupCodecId.NO_OP
+
+    override val name
+        get() = "N/A"
+
+    override fun encode(outputStream: OutputStream) = outputStream
+
+    override fun decode(inputStream: InputStream) = inputStream
+}
+
+/** Codec with ZIP compression. */
+class BackupZipCodec(
+    private val compressionLevel: Int,
+    override val name: String,
+) : BackupCodec {
+    override val id
+        get() = BackupCodecId.ZIP
+
+    override fun encode(outputStream: OutputStream) =
+        DeflaterOutputStream(outputStream, Deflater(compressionLevel))
+
+    override fun decode(inputStream: InputStream) = InflaterInputStream(inputStream)
+
+    companion object {
+        val DEFAULT_COMPRESSION = BackupZipCodec(Deflater.DEFAULT_COMPRESSION, "ZipDefault")
+        val BEST_COMPRESSION = BackupZipCodec(Deflater.BEST_COMPRESSION, "ZipBestCompression")
+        val BEST_SPEED = BackupZipCodec(Deflater.BEST_SPEED, "ZipBestSpeed")
+    }
+}
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreEntity.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreEntity.kt
index 6a7ef5a..817ee4c 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreEntity.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreEntity.kt
@@ -36,6 +36,13 @@
     val key: String
 
     /**
+     * Codec used to encode/decode the backup data.
+     *
+     * When it is null, the [BackupRestoreStorage.defaultCodec] will be used.
+     */
+    fun codec(): BackupCodec? = null
+
+    /**
      * Backs up the entity.
      *
      * @param backupContext context for backup
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreFileArchiver.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreFileArchiver.kt
index a447ace..3db46509 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreFileArchiver.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreFileArchiver.kt
@@ -41,6 +41,11 @@
     override fun createBackupRestoreEntities(): List<BackupRestoreEntity> =
         fileStorages.map { it.toBackupRestoreEntity() }
 
+    override fun wrapBackupOutputStream(codec: BackupCodec, outputStream: OutputStream) =
+        outputStream
+
+    override fun wrapRestoreInputStream(codec: BackupCodec, inputStream: InputStream) = inputStream
+
     override fun restoreEntity(data: BackupDataInputStream) {
         val key = data.key
         val fileStorage = fileStorages.firstOrNull { it.storageFilePath == key }
@@ -56,11 +61,19 @@
                 File(context.dataDirCompat, key)
             }
         Log.i(LOG_TAG, "[$name] Restore ${data.size()} bytes for $key to $file")
+        val inputStream = LimitedNoCloseInputStream(data)
         try {
+            val codec = BackupCodec.fromId(inputStream.read().toByte())
+            if (fileStorage != null && fileStorage.defaultCodec().id != codec.id) {
+                Log.i(
+                    LOG_TAG,
+                    "[$name] $key different codec: ${codec.id}, ${fileStorage.defaultCodec().id}"
+                )
+            }
             file.parentFile?.mkdirs() // ensure parent folders are created
-            val wrappedInputStream = wrapRestoreInputStream(data)
-            file.outputStream().use { wrappedInputStream.copyTo(it) }
-            Log.i(LOG_TAG, "[$name] $key restored")
+            val wrappedInputStream = codec.decode(inputStream)
+            val bytesCopied = file.outputStream().use { wrappedInputStream.copyTo(it) }
+            Log.i(LOG_TAG, "[$name] $key restore $bytesCopied bytes with ${codec.name}")
             fileStorage?.onRestoreFinished(file)
         } catch (e: Exception) {
             Log.e(LOG_TAG, "[$name] Fail to restore $key", e)
@@ -90,8 +103,12 @@
                 Log.i(LOG_TAG, "[$name] $key not exist")
                 return EntityBackupResult.DELETE
             }
-            val wrappedOutputStream = wrapBackupOutputStream(outputStream)
-            file.inputStream().use { it.copyTo(wrappedOutputStream) }
+            val codec = codec() ?: defaultCodec()
+            // MUST close to flush the data
+            wrapBackupOutputStream(codec, outputStream).use { stream ->
+                val bytesCopied = file.inputStream().use { it.copyTo(stream) }
+                Log.i(LOG_TAG, "[$name] $key backup $bytesCopied bytes with ${codec.name}")
+            }
             onBackupFinished(file)
             return EntityBackupResult.UPDATE
         }
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorage.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorage.kt
index 88d9dd6..8ff4bc8 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorage.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorage.kt
@@ -50,6 +50,9 @@
     /** Entities to back up and restore. */
     abstract fun createBackupRestoreEntities(): List<BackupRestoreEntity>
 
+    /** Default codec used to encode/decode the entity data. */
+    open fun defaultCodec(): BackupCodec = BackupZipCodec.BEST_COMPRESSION
+
     override fun performBackup(
         oldState: ParcelFileDescriptor?,
         data: BackupDataOutput,
@@ -64,9 +67,10 @@
         for (entity in entities) {
             val key = entity.key
             val outputStream = ByteArrayOutputStream()
+            val codec = entity.codec() ?: defaultCodec()
             val result =
                 try {
-                    entity.backup(backupContext, wrapBackupOutputStream(outputStream))
+                    entity.backup(backupContext, wrapBackupOutputStream(codec, outputStream))
                 } catch (exception: Exception) {
                     Log.e(LOG_TAG, "[$name] Fail to backup entity $key", exception)
                     continue
@@ -94,8 +98,10 @@
     /** Returns if backup is enabled. */
     open fun enableBackup(backupContext: BackupContext): Boolean = true
 
-    fun wrapBackupOutputStream(outputStream: OutputStream): OutputStream {
-        return outputStream
+    open fun wrapBackupOutputStream(codec: BackupCodec, outputStream: OutputStream): OutputStream {
+        // write a codec id header for safe restore
+        outputStream.write(codec.id.toInt())
+        return codec.encode(outputStream)
     }
 
     override fun restoreEntity(data: BackupDataInputStream) {
@@ -111,8 +117,12 @@
         }
         Log.i(LOG_TAG, "[$name] Restore $key: ${data.size()} bytes")
         val restoreContext = RestoreContext(key)
+        val codec = entity.codec() ?: defaultCodec()
         try {
-            entity.restore(restoreContext, wrapRestoreInputStream(data))
+            entity.restore(
+                restoreContext,
+                wrapRestoreInputStream(codec, LimitedNoCloseInputStream(data))
+            )
         } catch (exception: Exception) {
             Log.e(LOG_TAG, "[$name] Fail to restore entity $key", exception)
         }
@@ -121,8 +131,16 @@
     /** Returns if restore is enabled. */
     open fun enableRestore(): Boolean = true
 
-    fun wrapRestoreInputStream(inputStream: BackupDataInputStream): InputStream {
-        return LimitedNoCloseInputStream(inputStream)
+    open fun wrapRestoreInputStream(
+        codec: BackupCodec,
+        inputStream: InputStream,
+    ): InputStream {
+        // read the codec id first to check if it is expected codec
+        val id = inputStream.read()
+        val expectedId = codec.id.toInt()
+        if (id == expectedId) return codec.decode(inputStream)
+        Log.i(LOG_TAG, "Expect codec id $expectedId but got $id")
+        return BackupCodec.fromId(id.toByte()).decode(inputStream)
     }
 
     override fun writeNewStateDescription(newState: ParcelFileDescriptor) {}