Encrypted file storage for data store

EncryptedFileStorage.kt

package com.a.b

import android.content.Context
import androidx.annotation.GuardedBy
import androidx.datastore.core.ReadScope
import androidx.datastore.core.Serializer
import androidx.datastore.core.Storage
import androidx.datastore.core.StorageConnection
import androidx.datastore.core.WriteScope
import androidx.datastore.core.use
import androidx.datastore.preferences.core.Preferences
import androidx.datastore.preferences.core.PreferencesSerializer
import androidx.security.crypto.EncryptedFile
import androidx.security.crypto.MasterKey
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import okio.buffer
import okio.sink
import okio.source
import java.io.File
import java.io.FileNotFoundException
import java.io.FileOutputStream
import java.io.IOException
import java.io.InputStream
import java.io.OutputStream
import java.util.concurrent.atomic.AtomicBoolean

class EncryptedFileStorage<T>(
    private val context: Context,
    private val serializer: Serializer<T>,
    private val produceFile: () -> File,
) : Storage<T> {
    override fun createConnection(): StorageConnection<T> {
        val file = produceFile()

        synchronized(activeFilesLock) {
            val path = file.absolutePath
            check(!activeFiles.contains(path)) {
                "There are multiple DataStores active for the same file: $path. You should " +
                        "either maintain your DataStore as a singleton or confirm that there is " +
                        "no two DataStore's active on the same file (by confirming that the scope" +
                        " is cancelled)."
            }
            activeFiles.add(path)
        }

        val key = MasterKey.Builder(context)
            .setKeyScheme(MasterKey.KeyScheme.AES256_GCM)
            .build()

        return EncryptedFileStorageConnection(context, file, key, serializer) {
            synchronized(activeFilesLock) {
                activeFiles.remove(file.absolutePath)
            }
        }
    }

    internal companion object {
        /**
         * Active files should contain the absolute path for which there are currently active
         * DataStores. A DataStore is active until the scope it was created with has been
         * cancelled. Files aren't added to this list until the first read/write because the file
         * path is computed asynchronously.
         */
        @GuardedBy("activeFilesLock")
        internal val activeFiles = mutableSetOf<String>()

        internal val activeFilesLock = Any()
    }
}

class EncryptedFileStorageConnection<T>(
    private val context: Context,
    private val backingFile: File,
    private val key: MasterKey,
    private val serializer: Serializer<T>,
    private val onClose: () -> Unit,
) : StorageConnection<T> {

    private val closed = AtomicBoolean(false)
    private val transactionMutex = Mutex()

    override suspend fun <R> readScope(
        block: suspend ReadScope<T>.(locked: Boolean) -> R,
    ): R {
        checkNotClosed()

        val lock = transactionMutex.tryLock()
        try {
            val encryptedFile = EncryptedFile.Builder(context, backingFile, key, EncryptedFile.FileEncryptionScheme.AES256_GCM_HKDF_4KB)
                .build()

            return EncryptedFileReadScope(backingFile, encryptedFile, serializer).use {
                block(it, lock)
            }
        } finally {
            if (lock) {
                transactionMutex.unlock()
            }
        }
    }

    override suspend fun writeScope(block: suspend WriteScope<T>.() -> Unit) {
        checkNotClosed()
        backingFile.parentFile?.mkdirs()

        transactionMutex.withLock {
            val scratchFile = File(backingFile.absolutePath + ".tmp")
            if (scratchFile.exists()) {
                scratchFile.delete()
            }
            try {
                if (backingFile.exists() && !backingFile.renameTo(scratchFile)) {
                    throw IOException("Failed to rename backing file to scratch file")
                }
                val newEncryptedFile = EncryptedFile.Builder(context, backingFile, key, EncryptedFile.FileEncryptionScheme.AES256_GCM_HKDF_4KB)
                    .build()
                EncryptedFileWriteScope(backingFile, newEncryptedFile, serializer).use {
                    block(it)
                }
                if (!backingFile.exists()) {
                    scratchFile.renameTo(backingFile)
                    throw IOException(
                        "Backing file failed to generate a new file. Re-using old file, changes will be lost!"
                    )
                } else {
                    scratchFile.delete() // File was correctly written, delete old scracth file
                }
            } catch (ex: IOException) {
                if (scratchFile.exists()) {
                    if (backingFile.exists()) {
                        backingFile.delete()
                    }
                    if (!scratchFile.renameTo(backingFile)) {
                        scratchFile.delete()
                        throw IOException("Failed to rename old scratch file to backing file, storage lost", ex)
                    }
                }
                throw ex
            }
        }
    }

    override fun close() {
        closed.set(true)
        onClose()
    }

    private fun checkNotClosed() {
        check(!closed.get()) { "StorageConnection has already been disposed." }
    }
}

private open class EncryptedFileReadScope<T>(
    protected val backingFile: File,
    protected val file: EncryptedFile,
    protected val serializer: Serializer<T>,
) : ReadScope<T> {

    protected val closed = AtomicBoolean(false)

    override suspend fun readData(): T {
        checkNotClosed()
        return try {
            file.openFileInput().use { stream ->
                serializer.readFrom(stream)
            }
        } catch (ignored: FileNotFoundException) {
            if (backingFile.exists()) {
                // Re-read to prevent throwing from a race condition where the file is created by
                // another process after the initial read attempt but before `file.exists()` is
                // called. Otherwise file exists but we can't read it; throw FileNotFoundException
                // because something is wrong.
                return file.openFileInput().use { stream ->
                    serializer.readFrom(stream)
                }
            }
            return serializer.defaultValue
        }
    }

    override fun close() {
        closed.set(true)
    }

    protected fun checkNotClosed() {
        check(!closed.get()) { "This scope has already been closed." }
    }
}


private class EncryptedFileWriteScope<T>(backingFile: File, encryptedFile: EncryptedFile, serializer: Serializer<T>) :
    EncryptedFileReadScope<T>(backingFile, encryptedFile, serializer), WriteScope<T> {

    override suspend fun writeData(value: T) {
        checkNotClosed()
        file.openFileOutput().use { stream ->
            serializer.writeTo(value, UncloseableOutputStream(stream))
            stream.flush()
            stream.fd.sync()
        }
    }
}

private class UncloseableOutputStream(val fileOutputStream: FileOutputStream) : OutputStream() {

    override fun write(b: Int) {
        fileOutputStream.write(b)
    }

    override fun write(b: ByteArray) {
        fileOutputStream.write(b)
    }

    override fun write(bytes: ByteArray, off: Int, len: Int) {
        fileOutputStream.write(bytes, off, len)
    }

    override fun close() {
        // We will not close the underlying FileOutputStream until after we're done syncing
        // the fd. This is useful for things like b/173037611.
    }

    override fun flush() {
        fileOutputStream.flush()
    }
}

class FilePreferencesSerializer : Serializer<Preferences> {

    private val delegate = PreferencesSerializer

    override val defaultValue: Preferences
        get() = delegate.defaultValue

    override suspend fun readFrom(input: InputStream): Preferences {
        return input.source().buffer().use { buffer ->
            delegate.readFrom(buffer)
        }
    }

    override suspend fun writeTo(t: Preferences, output: OutputStream) {
        output.sink().buffer().use { buffer ->
            delegate.writeTo(t, buffer)
            buffer.flush()
        }
    }

}

Hey, what imports are you using for this gist?
Tried "implementation 'androidx.datastore:datastore-core-android:1.1.0-alpha01'" and also without the alpha01, but with no success.

androidx.security:security-crypto-ktx:1.1.0-alpha04 and androidx.datastore:datastore-preferences:1.1.0-alpha01 (last one would pull the core dependency probably)

Looks like v1.1.0 is out. It'd be great if we could have something like this as a part or extension of the official datastore library.