Skip to content

Commit 2056cbc

Browse files
SessionHero01SessionHero01
authored
Add encryption stream support (#20)
1 parent 3abec4a commit 2056cbc

File tree

6 files changed

+422
-0
lines changed

6 files changed

+422
-0
lines changed

‎gradle/libs.versions.toml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,10 @@ agp = "8.12.0"
33
kotlin = "2.2.0"
44

55
[libraries]
6+
junit = { module = "junit:junit", version = "4.13.2" }
7+
androidx-test-runner = { module = "androidx.test:runner", version = "1.7.0" }
8+
androidx-test-rules = { module = "androidx.test:rules", version = "1.7.0" }
9+
androidx-test-ext = { module = "androidx.test.ext:junit", version = "1.3.0" }
610

711
[plugins]
812
android-library = { id = "com.android.library", version.ref = "agp" }

‎library/build.gradle.kts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,4 +94,8 @@ publishing {
9494
}
9595

9696
dependencies {
97+
androidTestImplementation(libs.junit)
98+
androidTestImplementation(libs.androidx.test.runner)
99+
androidTestImplementation(libs.androidx.test.rules)
100+
androidTestImplementation(libs.androidx.test.ext)
97101
}

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

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
package network.loki.messenger.libsession_util
2+
3+
import androidx.test.ext.junit.runners.AndroidJUnit4
4+
import network.loki.messenger.libsession_util.encrypt.DecryptionStream
5+
import network.loki.messenger.libsession_util.encrypt.EncryptionStream
6+
import org.junit.Assert.assertArrayEquals
7+
import org.junit.Test
8+
import org.junit.runner.RunWith
9+
import java.io.ByteArrayInputStream
10+
import java.io.ByteArrayOutputStream
11+
import java.security.SecureRandom
12+
13+
@RunWith(AndroidJUnit4::class)
14+
class EncryptionStreamTest {
15+
16+
private fun testEncryptionCase(
17+
chunkSize: Int,
18+
dataSize: Int,
19+
) {
20+
try {
21+
val key = ByteArray(32)
22+
SecureRandom().nextBytes(key)
23+
24+
val expectData = ByteArray(dataSize)
25+
SecureRandom().nextBytes(expectData)
26+
27+
val encrypted = ByteArrayOutputStream().let { outputStream ->
28+
EncryptionStream(outputStream, key, chunkSize).use {
29+
it.write(expectData)
30+
}
31+
32+
outputStream.toByteArray()
33+
}
34+
35+
val actualData = DecryptionStream(
36+
ByteArrayInputStream(encrypted),
37+
key
38+
).use { it.readAllBytes() }
39+
40+
assertArrayEquals(expectData, actualData)
41+
} catch (e: Exception) {
42+
throw RuntimeException("Encryption/Decryption failed for chunkSize: $chunkSize, dataSize: $dataSize", e)
43+
}
44+
}
45+
46+
@Test
47+
fun shouldEncryptDecrypt() {
48+
testEncryptionCase(chunkSize = 24, dataSize = 25)
49+
testEncryptionCase(chunkSize = 24, dataSize = 24)
50+
testEncryptionCase(chunkSize = 24, dataSize = 12)
51+
testEncryptionCase(chunkSize = 24, dataSize = 48)
52+
}
53+
}

‎library/src/main/cpp/encryption.cpp‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,4 +156,134 @@ Java_network_loki_messenger_libsession_1util_SessionEncrypt_calculateECHDAgreeme
156156
return util::bytes_from_span(env, shared_secret);
157157
});
158158

159+
}
160+
161+
extern "C"
162+
JNIEXPORT jlong JNICALL
163+
Java_network_loki_messenger_libsession_1util_encrypt_EncryptionStream_00024Companion_createEncryptionStreamState(
164+
JNIEnv *env, jobject thiz, jbyteArray javaKey, jbyteArray javaHeaderOut) {
165+
JavaByteArrayRef key(env, javaKey);
166+
JavaByteArrayRef headerOut(env, javaHeaderOut);
167+
168+
if (headerOut.get().size() < crypto_secretstream_xchacha20poly1305_HEADERBYTES) {
169+
env->ThrowNew(env->FindClass("java/lang/IllegalArgumentException"),
170+
"Invalid headerOut: not enough space");
171+
return 0;
172+
}
173+
174+
if (key.get().size() != crypto_secretstream_xchacha20poly1305_KEYBYTES) {
175+
env->ThrowNew(env->FindClass("java/lang/IllegalArgumentException"),
176+
"Invalid key: unexpected size");
177+
return 0;
178+
}
179+
180+
auto state = std::make_unique<crypto_secretstream_xchacha20poly1305_state>();
181+
crypto_secretstream_xchacha20poly1305_init_push(state.get(),
182+
headerOut.get().data(),
183+
key.get().data());
184+
185+
return reinterpret_cast<jlong>(state.release());
186+
}
187+
188+
extern "C"
189+
JNIEXPORT jint JNICALL
190+
Java_network_loki_messenger_libsession_1util_encrypt_EncryptionStream_00024Companion_encryptionStreamHeaderSize(
191+
JNIEnv *env, jobject thiz) {
192+
return static_cast<jint>(crypto_secretstream_xchacha20poly1305_HEADERBYTES);
193+
}
194+
195+
extern "C"
196+
JNIEXPORT jint JNICALL
197+
Java_network_loki_messenger_libsession_1util_encrypt_EncryptionStream_00024Companion_encryptionStreamChunkOverhead(
198+
JNIEnv *env, jobject thiz) {
199+
return static_cast<jint>(crypto_secretstream_xchacha20poly1305_ABYTES);
200+
}
201+
202+
extern "C"
203+
JNIEXPORT jint JNICALL
204+
Java_network_loki_messenger_libsession_1util_encrypt_EncryptionStream_00024Companion_encryptStreamPush(
205+
JNIEnv *env, jobject thiz, jlong state_ptr, jbyteArray java_in_buf, jint in_buf_size, jbyteArray java_out_buf) {
206+
auto state = reinterpret_cast<crypto_secretstream_xchacha20poly1305_state*>(state_ptr);
207+
208+
JavaByteArrayRef in_buf(env, java_in_buf);
209+
JavaByteArrayRef out_buf(env, java_out_buf);
210+
211+
unsigned long long cipher_len = out_buf.get().size();
212+
213+
if (crypto_secretstream_xchacha20poly1305_push(
214+
state,
215+
out_buf.get().data(), &cipher_len, // Cipher data out
216+
in_buf.get().data(), in_buf_size, // Plaintext data in
217+
nullptr, 0, // Additional data (not used here)
218+
0 // Tag (not used here, can be 0 for message)
219+
)) {
220+
env->ThrowNew(env->FindClass("java/lang/IllegalStateException"),
221+
"Failed to push data into encryption stream");
222+
return 0;
223+
}
224+
225+
// Return the size of the ciphertext written to the output buffer
226+
return cipher_len;
227+
}
228+
229+
extern "C"
230+
JNIEXPORT void JNICALL
231+
Java_network_loki_messenger_libsession_1util_encrypt_EncryptionStream_00024Companion_destroyEncryptionStreamState(
232+
JNIEnv *env, jobject thiz, jlong state_ptr) {
233+
delete reinterpret_cast<crypto_secretstream_xchacha20poly1305_state*>(state_ptr);
234+
}
235+
236+
237+
extern "C"
238+
JNIEXPORT jlong JNICALL
239+
Java_network_loki_messenger_libsession_1util_encrypt_DecryptionStream_00024Companion_createDecryptionStreamState(
240+
JNIEnv *env, jobject thiz, jbyteArray javaKey, jbyteArray javaHeader) {
241+
JavaByteArrayRef key(env, javaKey);
242+
JavaByteArrayRef header(env, javaHeader);
243+
244+
if (header.get().size() < crypto_secretstream_xchacha20poly1305_HEADERBYTES) {
245+
env->ThrowNew(env->FindClass("java/lang/IllegalArgumentException"),
246+
"Invalid header: unexpected size");
247+
return 0;
248+
}
249+
250+
if (key.get().size() != crypto_secretstream_xchacha20poly1305_KEYBYTES) {
251+
env->ThrowNew(env->FindClass("java/lang/IllegalArgumentException"),
252+
"Invalid key: unexpected size");
253+
return 0;
254+
}
255+
256+
auto state = std::make_unique<crypto_secretstream_xchacha20poly1305_state>();
257+
258+
if (crypto_secretstream_xchacha20poly1305_init_pull(state.get(), header.get().data(), key.get().data()) != 0) {
259+
env->ThrowNew(env->FindClass("java/lang/IllegalArgumentException"),
260+
"Failed to initialize decryption stream state");
261+
return 0;
262+
}
263+
264+
return reinterpret_cast<jlong>(state.release());
265+
}
266+
267+
extern "C"
268+
JNIEXPORT jint JNICALL
269+
Java_network_loki_messenger_libsession_1util_encrypt_DecryptionStream_00024Companion_decryptionStreamPull(
270+
JNIEnv *env, jobject thiz, jlong native_state_ptr, jbyteArray java_in_buf, jint in_buf_len, jbyteArray java_out_buf) {
271+
JavaByteArrayRef out_buf(env, java_out_buf);
272+
JavaByteArrayRef in_buf(env, java_in_buf);
273+
274+
unsigned long long mlen = out_buf.get().size();
275+
276+
if (crypto_secretstream_xchacha20poly1305_pull(
277+
reinterpret_cast<crypto_secretstream_xchacha20poly1305_state*>(native_state_ptr),
278+
out_buf.get().data(), &mlen, // Plaintext data out
279+
nullptr, // Tag (not used here)
280+
in_buf.get().data(), in_buf_len, // Ciphertext data in
281+
nullptr, 0 // Additional data (not used here)
282+
)) {
283+
env->ThrowNew(env->FindClass("java/lang/IllegalStateException"),
284+
"Failed to pull data from decryption stream");
285+
return 0;
286+
}
287+
288+
return mlen;
159289
}

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

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
package network.loki.messenger.libsession_util.encrypt
2+
3+
import java.io.InputStream
4+
import java.nio.ByteBuffer
5+
6+
/**
7+
* A stream for decrypting data that was encrypted using [EncryptionStream].
8+
*
9+
* Note that you don't need to specify the chunk size as it would have been encoded into the stream.
10+
*
11+
* @param inStream the underlying input stream to read encrypted data from.
12+
* @param key the decryption key to use for decrypting data, must be 32 bytes long.
13+
* @param autoClose whether to automatically close the underlying input stream when this stream is
14+
*/
15+
class DecryptionStream(
16+
private val inStream: InputStream,
17+
key: ByteArray,
18+
private val autoClose: Boolean = true
19+
) : InputStream() {
20+
private val nativeStatePtr: Long
21+
private val chunkSize: Int
22+
private val cipherBuffer: ByteArray
23+
private val plaintextBuffer: ByteBuffer
24+
25+
init {
26+
try {
27+
// Read chunk size from the input stream
28+
val chunkSizeBuffer = ByteBuffer.allocate(4)
29+
check(inStream.read(chunkSizeBuffer.array(), 0, 4) == 4) {
30+
"Failed to read the chunk size from the input stream."
31+
}
32+
33+
chunkSize = chunkSizeBuffer.int
34+
35+
cipherBuffer = ByteArray((chunkSize + EncryptionStream.encryptionStreamChunkOverhead()).coerceAtLeast(
36+
EncryptionStream.encryptionStreamHeaderSize()
37+
))
38+
39+
plaintextBuffer = ByteBuffer.allocate(chunkSize)
40+
plaintextBuffer.limit(0) // Should be empty initially
41+
42+
// Read the initial header from the input stream
43+
check(inStream.read(cipherBuffer, 0, EncryptionStream.encryptionStreamHeaderSize())
44+
== EncryptionStream.encryptionStreamHeaderSize()) {
45+
"Failed to read the initial header from the input stream."
46+
}
47+
48+
// Initialize the native decryption stream state
49+
nativeStatePtr = createDecryptionStreamState(key, cipherBuffer)
50+
} catch (e: Exception) {
51+
if (autoClose) {
52+
inStream.close()
53+
}
54+
55+
throw e
56+
}
57+
}
58+
59+
protected fun finalize() {
60+
EncryptionStream.destroyEncryptionStreamState(nativeStatePtr)
61+
}
62+
63+
private fun readChunkIfNeeded() {
64+
if (!plaintextBuffer.hasRemaining()) {
65+
// Read the next chunk of encrypted data from the input stream
66+
val bytesRead = inStream.read(cipherBuffer)
67+
if (bytesRead == -1) {
68+
// End of stream reached
69+
return
70+
}
71+
72+
// Decrypt the chunk and fill the plaintext buffer
73+
plaintextBuffer.clear()
74+
val decryptedBytes = decryptionStreamPull(nativeStatePtr, cipherBuffer, bytesRead, plaintextBuffer.array())
75+
if (decryptedBytes < 0) {
76+
throw IllegalStateException("Decryption failed with error code: $decryptedBytes")
77+
}
78+
79+
plaintextBuffer.limit(decryptedBytes)
80+
}
81+
}
82+
83+
override fun read(): Int {
84+
readChunkIfNeeded()
85+
86+
if (plaintextBuffer.hasRemaining()) {
87+
return plaintextBuffer.get().toInt()
88+
}
89+
90+
return -1 // End of stream
91+
}
92+
93+
override fun read(b: ByteArray, off: Int, len: Int): Int {
94+
readChunkIfNeeded()
95+
96+
if (!plaintextBuffer.hasRemaining()) {
97+
return -1 // End of stream
98+
}
99+
100+
val bytesToRead = minOf(len, plaintextBuffer.remaining())
101+
plaintextBuffer.get(b, off, bytesToRead)
102+
return bytesToRead
103+
}
104+
105+
override fun close() {
106+
if (autoClose) {
107+
inStream.close()
108+
}
109+
110+
super.close()
111+
}
112+
113+
companion object {
114+
private external fun createDecryptionStreamState(key: ByteArray, header: ByteArray): Long
115+
private external fun decryptionStreamPull(nativeStatePtr: Long, inBuf: ByteArray, inLen: Int, outBuf: ByteArray): Int
116+
}
117+
}

0 commit comments

Comments
 (0)