Skip to content

Commit 795aef8

Browse files
SessionHero01SessionHero01
SessionHero01
authored and
SessionHero01
committed
Encrypt streaming support
1 parent 3abec4a commit 795aef8

File tree

3 files changed

+336
-0
lines changed

3 files changed

+336
-0
lines changed

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

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,4 +156,135 @@ 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+
reinterpret_cast<unsigned char*>(env->GetDirectBufferAddress(javaHeaderOut)),
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+
unsigned char tag;
276+
277+
if (crypto_secretstream_xchacha20poly1305_pull(
278+
reinterpret_cast<crypto_secretstream_xchacha20poly1305_state*>(native_state_ptr),
279+
out_buf.get().data(), &mlen, // Plaintext data out
280+
&tag,
281+
in_buf.get().data(), in_buf_len, // Ciphertext data in
282+
nullptr, 0 // Additional data (not used here)
283+
)) {
284+
env->ThrowNew(env->FindClass("java/lang/IllegalStateException"),
285+
"Failed to pull data from decryption stream");
286+
return 0;
287+
}
288+
289+
return mlen;
159290
}

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

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

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

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
package network.loki.messenger.libsession_util.encrypt
2+
3+
import java.io.OutputStream
4+
import java.nio.ByteBuffer
5+
6+
class EncryptionStream(
7+
private val out: OutputStream,
8+
key: ByteArray,
9+
val chunkSize: Int = 4096,
10+
private val autoClose: Boolean = true
11+
) : OutputStream() {
12+
private val nativeStatePtr: Long
13+
private val plaintextBuffer by lazy(LazyThreadSafetyMode.NONE) {
14+
ByteBuffer.allocate(chunkSize)
15+
}
16+
17+
private val cipherBuffer = ByteBuffer.allocate(chunkSize + encryptionStreamChunkOverhead())
18+
19+
init {
20+
try {
21+
// Write the chunk size
22+
out.write(ByteBuffer.allocate(4).putInt(chunkSize).array())
23+
24+
// Write initial header to the output stream
25+
nativeStatePtr = createEncryptionStreamState(key, cipherBuffer.array())
26+
out.write(cipherBuffer.array(), 0, encryptionStreamHeaderSize())
27+
} catch (e: Exception) {
28+
if (autoClose) {
29+
out.close()
30+
}
31+
32+
throw e
33+
}
34+
}
35+
36+
protected fun finalize() {
37+
destroyEncryptionStreamState(nativeStatePtr)
38+
}
39+
40+
private fun flushWhenFull() {
41+
if (!plaintextBuffer.hasRemaining()) {
42+
flush()
43+
}
44+
}
45+
46+
override fun write(b: Int) {
47+
flushWhenFull()
48+
plaintextBuffer.put(b.toByte())
49+
flushWhenFull()
50+
}
51+
52+
override fun write(b: ByteArray, off: Int, len: Int) {
53+
var from = off
54+
val to = off + len
55+
while (from < to) {
56+
flushWhenFull()
57+
val toWrite = minOf(to - from, chunkSize - plaintextBuffer.position()).coerceAtLeast(0)
58+
plaintextBuffer.put(b, from, toWrite)
59+
from += toWrite
60+
flushWhenFull()
61+
}
62+
}
63+
64+
override fun flush() {
65+
super.flush()
66+
67+
// Flip the buffer to prepare for reading
68+
plaintextBuffer.flip()
69+
70+
if (plaintextBuffer.hasRemaining()) {
71+
cipherBuffer.clear()
72+
val cipherLength = encryptStreamPush(nativeStatePtr, plaintextBuffer.array(), plaintextBuffer.remaining(), cipherBuffer.array())
73+
out.write(cipherBuffer.array(), 0, cipherLength)
74+
plaintextBuffer.clear()
75+
}
76+
77+
out.flush()
78+
}
79+
80+
override fun close() {
81+
flush()
82+
if (autoClose) {
83+
out.close()
84+
}
85+
super.close()
86+
}
87+
88+
89+
companion object {
90+
external fun encryptionStreamHeaderSize(): Int
91+
external fun encryptionStreamChunkOverhead(): Int
92+
93+
private external fun createEncryptionStreamState(key: ByteArray, headerOut: ByteArray): Long
94+
private external fun encryptStreamPush(statePtr: Long, inBuf: ByteArray, inBufLen: Int, outBuf: ByteArray): Int
95+
96+
external fun destroyEncryptionStreamState(statePtr: Long)
97+
}
98+
}

0 commit comments

Comments
 (0)