Add encryption stream support

gradle/libs.versions.toml

@@ -3,6 +3,10 @@ agp = "8.12.0"
 kotlin = "2.2.0"
 
 [libraries]
+junit = { module = "junit:junit", version = "4.13.2" }
+androidx-test-runner = { module = "androidx.test:runner", version = "1.7.0" }
+androidx-test-rules = { module = "androidx.test:rules", version = "1.7.0" }
+androidx-test-ext = { module = "androidx.test.ext:junit", version = "1.3.0" }
 
 [plugins]
 android-library = { id = "com.android.library", version.ref = "agp" }

library/build.gradle.kts

@@ -94,4 +94,8 @@ publishing {
 }
 
 dependencies {
+    androidTestImplementation(libs.junit)
+    androidTestImplementation(libs.androidx.test.runner)
+    androidTestImplementation(libs.androidx.test.rules)
+    androidTestImplementation(libs.androidx.test.ext)
 }

library/src/androidTest/kotlin/network/loki/messenger/libsession_util/EncryptionStreamTest.kt

@@ -0,0 +1,53 @@
+package network.loki.messenger.libsession_util
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import network.loki.messenger.libsession_util.encrypt.DecryptionStream
+import network.loki.messenger.libsession_util.encrypt.EncryptionStream
+import org.junit.Assert.assertArrayEquals
+import org.junit.Test
+import org.junit.runner.RunWith
+import java.io.ByteArrayInputStream
+import java.io.ByteArrayOutputStream
+import java.security.SecureRandom
+
+@RunWith(AndroidJUnit4::class)
+class EncryptionStreamTest {
+
+    private fun testEncryptionCase(
+        chunkSize: Int,
+        dataSize: Int,
+    ) {
+        try {
+            val key = ByteArray(32)
+            SecureRandom().nextBytes(key)
+
+            val expectData = ByteArray(dataSize)
+            SecureRandom().nextBytes(expectData)
+
+            val encrypted = ByteArrayOutputStream().let { outputStream ->
+                EncryptionStream(outputStream, key, chunkSize).use {
+                    it.write(expectData)
+                }
+
+                outputStream.toByteArray()
+            }
+
+            val actualData = DecryptionStream(
+                ByteArrayInputStream(encrypted),
+                key
+            ).use { it.readAllBytes() }
+
+            assertArrayEquals(expectData, actualData)
+        } catch (e: Exception) {
+            throw RuntimeException("Encryption/Decryption failed for chunkSize: $chunkSize, dataSize: $dataSize", e)
+        }
+    }
+
+    @Test
+    fun shouldEncryptDecrypt() {
+        testEncryptionCase(chunkSize = 24, dataSize = 25)
+        testEncryptionCase(chunkSize = 24, dataSize = 24)
+        testEncryptionCase(chunkSize = 24, dataSize = 12)
+        testEncryptionCase(chunkSize = 24, dataSize = 48)
+    }
+}

library/src/main/cpp/encryption.cpp

@@ -156,4 +156,134 @@ Java_network_loki_messenger_libsession_1util_SessionEncrypt_calculateECHDAgreeme
         return util::bytes_from_span(env, shared_secret);
     });
 
+}
+
+extern "C"
+JNIEXPORT jlong JNICALL
+Java_network_loki_messenger_libsession_1util_encrypt_EncryptionStream_00024Companion_createEncryptionStreamState(
+        JNIEnv *env, jobject thiz, jbyteArray javaKey, jbyteArray javaHeaderOut) {
+    JavaByteArrayRef key(env, javaKey);
+    JavaByteArrayRef headerOut(env, javaHeaderOut);
+
+    if (headerOut.get().size() < crypto_secretstream_xchacha20poly1305_HEADERBYTES) {
+        env->ThrowNew(env->FindClass("java/lang/IllegalArgumentException"),
+                      "Invalid headerOut: not enough space");
+        return 0;
+    }
+
+    if (key.get().size() != crypto_secretstream_xchacha20poly1305_KEYBYTES) {
+        env->ThrowNew(env->FindClass("java/lang/IllegalArgumentException"),
+                      "Invalid key: unexpected size");
+        return 0;
+    }
+
+    auto state = std::make_unique<crypto_secretstream_xchacha20poly1305_state>();
+    crypto_secretstream_xchacha20poly1305_init_push(state.get(),
+                                                    headerOut.get().data(),
+                                                    key.get().data());
+
+    return reinterpret_cast<jlong>(state.release());
+}
+
+extern "C"
+JNIEXPORT jint JNICALL
+Java_network_loki_messenger_libsession_1util_encrypt_EncryptionStream_00024Companion_encryptionStreamHeaderSize(
+        JNIEnv *env, jobject thiz) {
+    return static_cast<jint>(crypto_secretstream_xchacha20poly1305_HEADERBYTES);
+}
+
+extern "C"
+JNIEXPORT jint JNICALL
+Java_network_loki_messenger_libsession_1util_encrypt_EncryptionStream_00024Companion_encryptionStreamChunkOverhead(
+        JNIEnv *env, jobject thiz) {
+    return static_cast<jint>(crypto_secretstream_xchacha20poly1305_ABYTES);
+}
+
+extern "C"
+JNIEXPORT jint JNICALL
+Java_network_loki_messenger_libsession_1util_encrypt_EncryptionStream_00024Companion_encryptStreamPush(
+        JNIEnv *env, jobject thiz, jlong state_ptr, jbyteArray java_in_buf, jint in_buf_size, jbyteArray java_out_buf) {
+    auto state = reinterpret_cast<crypto_secretstream_xchacha20poly1305_state*>(state_ptr);
+
+    JavaByteArrayRef in_buf(env, java_in_buf);
+    JavaByteArrayRef out_buf(env, java_out_buf);
+
+    unsigned long long cipher_len = out_buf.get().size();
+
+    if (crypto_secretstream_xchacha20poly1305_push(
+            state,
+            out_buf.get().data(), &cipher_len, // Cipher data out
+            in_buf.get().data(), in_buf_size, // Plaintext data in
+            nullptr, 0, // Additional data (not used here)
+            0 // Tag (not used here, can be 0 for message)
+    )) {
+        env->ThrowNew(env->FindClass("java/lang/IllegalStateException"),
+                      "Failed to push data into encryption stream");
+        return 0;
+    }
+
+    // Return the size of the ciphertext written to the output buffer
+    return cipher_len;
+}
+
+extern "C"
+JNIEXPORT void JNICALL
+Java_network_loki_messenger_libsession_1util_encrypt_EncryptionStream_00024Companion_destroyEncryptionStreamState(
+        JNIEnv *env, jobject thiz, jlong state_ptr) {
+    delete reinterpret_cast<crypto_secretstream_xchacha20poly1305_state*>(state_ptr);
+}
+
+
+extern "C"
+JNIEXPORT jlong JNICALL
+Java_network_loki_messenger_libsession_1util_encrypt_DecryptionStream_00024Companion_createDecryptionStreamState(
+        JNIEnv *env, jobject thiz, jbyteArray javaKey, jbyteArray javaHeader) {
+    JavaByteArrayRef key(env, javaKey);
+    JavaByteArrayRef header(env, javaHeader);
+
+    if (header.get().size() < crypto_secretstream_xchacha20poly1305_HEADERBYTES) {
+        env->ThrowNew(env->FindClass("java/lang/IllegalArgumentException"),
+                      "Invalid header: unexpected size");
+        return 0;
+    }
+
+    if (key.get().size() != crypto_secretstream_xchacha20poly1305_KEYBYTES) {
+        env->ThrowNew(env->FindClass("java/lang/IllegalArgumentException"),
+                      "Invalid key: unexpected size");
+        return 0;
+    }
+
+    auto state = std::make_unique<crypto_secretstream_xchacha20poly1305_state>();
+
+    if (crypto_secretstream_xchacha20poly1305_init_pull(state.get(), header.get().data(), key.get().data()) != 0) {
+        env->ThrowNew(env->FindClass("java/lang/IllegalArgumentException"),
+                      "Failed to initialize decryption stream state");
+        return 0;
+    }
+
+    return reinterpret_cast<jlong>(state.release());
+}
+
+extern "C"
+JNIEXPORT jint JNICALL
+Java_network_loki_messenger_libsession_1util_encrypt_DecryptionStream_00024Companion_decryptionStreamPull(
+        JNIEnv *env, jobject thiz, jlong native_state_ptr, jbyteArray java_in_buf, jint in_buf_len, jbyteArray java_out_buf) {
+    JavaByteArrayRef out_buf(env, java_out_buf);
+    JavaByteArrayRef in_buf(env, java_in_buf);
+
+    unsigned long long mlen = out_buf.get().size();
+
+    if (crypto_secretstream_xchacha20poly1305_pull(
+            reinterpret_cast<crypto_secretstream_xchacha20poly1305_state*>(native_state_ptr),
+            out_buf.get().data(), &mlen, // Plaintext data out
+            nullptr, // Tag (not used here)
+            in_buf.get().data(), in_buf_len, // Ciphertext data in
+            nullptr, 0 // Additional data (not used here)
+            )) {
+        env->ThrowNew(env->FindClass("java/lang/IllegalStateException"),
+                      "Failed to pull data from decryption stream");
+        return 0;
+    }
+
+    return mlen;
 }

library/src/main/java/network/loki/messenger/libsession_util/encrypt/DecryptionStream.kt

@@ -0,0 +1,117 @@
+package network.loki.messenger.libsession_util.encrypt
+
+import java.io.InputStream
+import java.nio.ByteBuffer
+
+/**
+ * A stream for decrypting data that was encrypted using [EncryptionStream].
+ *
+ * Note that you don't need to specify the chunk size as it would have been encoded into the stream.
+ *
+ * @param inStream the underlying input stream to read encrypted data from.
+ * @param key the decryption key to use for decrypting data, must be 32 bytes long.
+ * @param autoClose whether to automatically close the underlying input stream when this stream is
+ */
+class DecryptionStream(
+    private val inStream: InputStream,
+    key: ByteArray,
+    private val autoClose: Boolean = true
+) : InputStream() {
+    private val nativeStatePtr: Long
+    private val chunkSize: Int
+    private val cipherBuffer: ByteArray
+    private val plaintextBuffer: ByteBuffer
+
+    init {
+        try {
+            // Read chunk size from the input stream
+            val chunkSizeBuffer = ByteBuffer.allocate(4)
+            check(inStream.read(chunkSizeBuffer.array(), 0, 4) == 4) {
+                "Failed to read the chunk size from the input stream."
+            }
+
+            chunkSize = chunkSizeBuffer.int
+
+            cipherBuffer = ByteArray((chunkSize + EncryptionStream.encryptionStreamChunkOverhead()).coerceAtLeast(
+                EncryptionStream.encryptionStreamHeaderSize()
+            ))
+
+            plaintextBuffer = ByteBuffer.allocate(chunkSize)
+            plaintextBuffer.limit(0) // Should be empty initially
+
+            // Read the initial header from the input stream
+            check(inStream.read(cipherBuffer, 0, EncryptionStream.encryptionStreamHeaderSize())
+                == EncryptionStream.encryptionStreamHeaderSize()) {
+                "Failed to read the initial header from the input stream."
+            }
+
+            // Initialize the native decryption stream state
+            nativeStatePtr = createDecryptionStreamState(key, cipherBuffer)
+        } catch (e: Exception) {
+            if (autoClose) {
+                inStream.close()
+            }
+
+            throw e
+        }
+    }
+
+    protected fun finalize() {
+        EncryptionStream.destroyEncryptionStreamState(nativeStatePtr)
+    }
+
+    private fun readChunkIfNeeded() {
+        if (!plaintextBuffer.hasRemaining()) {
+            // Read the next chunk of encrypted data from the input stream
+            val bytesRead = inStream.read(cipherBuffer)
+            if (bytesRead == -1) {
+                // End of stream reached
+                return
+            }
+
+            // Decrypt the chunk and fill the plaintext buffer
+            plaintextBuffer.clear()
+            val decryptedBytes = decryptionStreamPull(nativeStatePtr, cipherBuffer, bytesRead, plaintextBuffer.array())
+            if (decryptedBytes < 0) {
+                throw IllegalStateException("Decryption failed with error code: $decryptedBytes")
+            }
+
+            plaintextBuffer.limit(decryptedBytes)
+        }
+    }
+
+    override fun read(): Int {
+        readChunkIfNeeded()
+
+        if (plaintextBuffer.hasRemaining()) {
+            return plaintextBuffer.get().toInt()
+        }
+
+        return -1 // End of stream
+    }
+
+    override fun read(b: ByteArray, off: Int, len: Int): Int {
+        readChunkIfNeeded()
+
+        if (!plaintextBuffer.hasRemaining()) {
+            return -1 // End of stream
+        }
+
+        val bytesToRead = minOf(len, plaintextBuffer.remaining())
+        plaintextBuffer.get(b, off, bytesToRead)
+        return bytesToRead
+    }
+
+    override fun close() {
+        if (autoClose) {
+            inStream.close()
+        }
+
+        super.close()
+    }
+
+    companion object {
+        private external fun createDecryptionStreamState(key: ByteArray, header: ByteArray): Long
+        private external fun decryptionStreamPull(nativeStatePtr: Long, inBuf: ByteArray, inLen: Int, outBuf: ByteArray): Int
+    }
+}