From f8d8558cdbf20f8e9946eb6739dde94116d76682 Mon Sep 17 00:00:00 2001 From: Greyson Parrelli Date: Wed, 11 Jun 2025 11:58:35 -0400 Subject: [PATCH] Convert AttachmentCipherTest to kotlin. --- .../signal/core/util/InputStreamExtensions.kt | 4 +- .../java/org/signal/core/util/StreamUtil.java | 9 +- .../api/crypto/AttachmentCipherTest.java | 535 ----------------- .../api/crypto/AttachmentCipherTest.kt | 548 ++++++++++++++++++ 4 files changed, 556 insertions(+), 540 deletions(-) delete mode 100644 libsignal-service/src/test/java/org/whispersystems/signalservice/api/crypto/AttachmentCipherTest.java create mode 100644 libsignal-service/src/test/java/org/whispersystems/signalservice/api/crypto/AttachmentCipherTest.kt diff --git a/core-util-jvm/src/main/java/org/signal/core/util/InputStreamExtensions.kt b/core-util-jvm/src/main/java/org/signal/core/util/InputStreamExtensions.kt index 645e48de72..3905def430 100644 --- a/core-util-jvm/src/main/java/org/signal/core/util/InputStreamExtensions.kt +++ b/core-util-jvm/src/main/java/org/signal/core/util/InputStreamExtensions.kt @@ -127,8 +127,8 @@ fun InputStream.limit(limit: Long): LimitedInputStream { * * @param closeInputStream If true, the input stream will be closed after the copy is complete. */ -fun InputStream.copyTo(outputStream: OutputStream, closeInputStream: Boolean = true): Long { - return StreamUtil.copy(this, outputStream, closeInputStream) +fun InputStream.copyTo(outputStream: OutputStream, closeInputStream: Boolean = true, closeOutputStream: Boolean = true): Long { + return StreamUtil.copy(this, outputStream, closeInputStream, closeOutputStream) } /** diff --git a/core-util-jvm/src/main/java/org/signal/core/util/StreamUtil.java b/core-util-jvm/src/main/java/org/signal/core/util/StreamUtil.java index 5c9c60617c..62d4b983d4 100644 --- a/core-util-jvm/src/main/java/org/signal/core/util/StreamUtil.java +++ b/core-util-jvm/src/main/java/org/signal/core/util/StreamUtil.java @@ -96,10 +96,10 @@ public final class StreamUtil { } public static long copy(InputStream in, OutputStream out) throws IOException { - return copy(in, out, true); + return copy(in, out, true, true); } - public static long copy(InputStream in, OutputStream out, boolean closeInputStream) throws IOException { + public static long copy(InputStream in, OutputStream out, boolean closeInputStream, boolean closeOutputStream) throws IOException { byte[] buffer = new byte[64 * 1024]; int read; long total = 0; @@ -114,7 +114,10 @@ public final class StreamUtil { } out.flush(); - out.close(); + + if (closeOutputStream) { + out.close(); + } return total; } diff --git a/libsignal-service/src/test/java/org/whispersystems/signalservice/api/crypto/AttachmentCipherTest.java b/libsignal-service/src/test/java/org/whispersystems/signalservice/api/crypto/AttachmentCipherTest.java deleted file mode 100644 index 8401843612..0000000000 --- a/libsignal-service/src/test/java/org/whispersystems/signalservice/api/crypto/AttachmentCipherTest.java +++ /dev/null @@ -1,535 +0,0 @@ -package org.whispersystems.signalservice.api.crypto; - -import org.conscrypt.Conscrypt; -import org.junit.Test; -import org.signal.core.util.StreamUtil; -import org.signal.libsignal.protocol.InvalidMessageException; -import org.signal.libsignal.protocol.incrementalmac.ChunkSizeChoice; -import org.signal.libsignal.protocol.incrementalmac.InvalidMacException; -import org.signal.libsignal.protocol.kdf.HKDF; -import org.whispersystems.signalservice.api.backup.MediaRootBackupKey; -import org.whispersystems.signalservice.internal.crypto.PaddingInputStream; -import org.whispersystems.signalservice.internal.push.http.AttachmentCipherOutputStreamFactory; -import org.whispersystems.signalservice.internal.util.Util; - -import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; -import java.io.File; -import java.io.FileOutputStream; -import java.io.IOException; -import java.io.InputStream; -import java.io.OutputStream; -import java.security.Security; -import java.util.Arrays; -import java.util.Random; - -import static org.junit.Assert.assertArrayEquals; -import static org.junit.Assert.assertTrue; -import static org.junit.Assert.fail; -import static org.whispersystems.signalservice.testutil.LibSignalLibraryUtil.assumeLibSignalSupportedOnOS; - -public final class AttachmentCipherTest { - - static { - // https://github.com/google/conscrypt/issues/1034 - if (!System.getProperty("os.arch").equals("aarch64")) { - Security.insertProviderAt(Conscrypt.newProvider(), 1); - } - } - - private static final int MEBIBYTE = 1024 * 1024; - - @Test - public void attachment_encryptDecrypt_nonIncremental() throws IOException, InvalidMessageException { - attachment_encryptDecrypt(false, MEBIBYTE); - } - - @Test - public void attachment_encryptDecrypt_incremental() throws IOException, InvalidMessageException { - attachment_encryptDecrypt(true, MEBIBYTE); - } - - @Test - public void attachment_encryptDecrypt_incremental_manyFileSizes() throws IOException, InvalidMessageException { - // Designed to stress the various boundary conditions of reading the final mac - for (int i = 0; i < 100; i++) { - attachment_encryptDecrypt(true, MEBIBYTE + new Random().nextInt(1, 64 * 1024)); - } - } - - private void attachment_encryptDecrypt(boolean incremental, int fileSize) throws IOException, InvalidMessageException { - byte[] key = Util.getSecretBytes(64); - byte[] plaintextInput = Util.getSecretBytes(fileSize); - EncryptResult encryptResult = encryptData(plaintextInput, key, incremental); - File cipherFile = writeToFile(encryptResult.ciphertext); - InputStream inputStream = AttachmentCipherInputStream.createForAttachment(cipherFile, plaintextInput.length, key, encryptResult.digest, encryptResult.incrementalDigest, encryptResult.chunkSizeChoice); - byte[] plaintextOutput = readInputStreamFully(inputStream); - - assertArrayEquals(plaintextInput, plaintextOutput); - - cipherFile.delete(); - } - - @Test - public void attachment_encryptDecryptEmpty_nonIncremental() throws IOException, InvalidMessageException { - attachment_encryptDecryptEmpty(false); - } - - @Test - public void attachment_encryptDecryptEmpty_incremental() throws IOException, InvalidMessageException { - attachment_encryptDecryptEmpty(true); - } - - private void attachment_encryptDecryptEmpty(boolean incremental) throws IOException, InvalidMessageException { - byte[] key = Util.getSecretBytes(64); - byte[] plaintextInput = "".getBytes(); - EncryptResult encryptResult = encryptData(plaintextInput, key, incremental); - File cipherFile = writeToFile(encryptResult.ciphertext); - InputStream inputStream = AttachmentCipherInputStream.createForAttachment(cipherFile, plaintextInput.length, key, encryptResult.digest, encryptResult.incrementalDigest, encryptResult.chunkSizeChoice); - byte[] plaintextOutput = readInputStreamFully(inputStream); - - assertArrayEquals(plaintextInput, plaintextOutput); - - cipherFile.delete(); - } - - @Test(expected = InvalidMessageException.class) - public void attachment_decryptFailOnBadKey_nonIncremental() throws IOException, InvalidMessageException { - attachment_decryptFailOnBadKey(false); - } - - @Test(expected = InvalidMessageException.class) - public void attachment_decryptFailOnBadKey_incremental() throws IOException, InvalidMessageException { - attachment_decryptFailOnBadKey(true); - } - - private void attachment_decryptFailOnBadKey(boolean incremental) throws IOException, InvalidMessageException { - File cipherFile = null; - - try { - byte[] key = Util.getSecretBytes(64); - byte[] plaintextInput = Util.getSecretBytes(MEBIBYTE); - EncryptResult encryptResult = encryptData(plaintextInput, key, incremental); - byte[] badKey = new byte[64]; - - cipherFile = writeToFile(encryptResult.ciphertext); - - AttachmentCipherInputStream.createForAttachment(cipherFile, plaintextInput.length, badKey, encryptResult.digest, null, 0); - } finally { - if (cipherFile != null) { - cipherFile.delete(); - } - } - } - @Test(expected = InvalidMessageException.class) - public void attachment_decryptFailOnBadMac_nonIncremental() throws IOException, InvalidMessageException { - attachment_decryptFailOnBadMac(false); - } - - @Test(expected = InvalidMessageException.class) - public void attachment_decryptFailOnBadMac_incremental() throws IOException, InvalidMessageException { - attachment_decryptFailOnBadMac(true); - } - - private void attachment_decryptFailOnBadMac(boolean incremental) throws IOException, InvalidMessageException { - File cipherFile = null; - - try { - byte[] key = Util.getSecretBytes(64); - byte[] plaintextInput = Util.getSecretBytes(MEBIBYTE); - EncryptResult encryptResult = encryptData(plaintextInput, key, incremental); - byte[] badMacCiphertext = Arrays.copyOf(encryptResult.ciphertext, encryptResult.ciphertext.length); - - badMacCiphertext[badMacCiphertext.length - 1] += 1; - - cipherFile = writeToFile(badMacCiphertext); - - InputStream stream = AttachmentCipherInputStream.createForAttachment(cipherFile, plaintextInput.length, key, encryptResult.digest, encryptResult.incrementalDigest, encryptResult.chunkSizeChoice); - - // In incremental mode, we'll only check the digest after reading the whole thing - if (incremental) { - StreamUtil.readFully(stream); - } - } finally { - if (cipherFile != null) { - cipherFile.delete(); - } - } - } - - @Test(expected = InvalidMessageException.class) - public void attachment_decryptFailOnNullDigest_nonIncremental() throws IOException, InvalidMessageException { - attachment_decryptFailOnNullDigest(false); - } - - @Test(expected = InvalidMessageException.class) - public void attachment_decryptFailOnNullDigest_incremental() throws IOException, InvalidMessageException { - attachment_decryptFailOnNullDigest(true); - } - - private void attachment_decryptFailOnNullDigest(boolean incremental) throws IOException, InvalidMessageException { - File cipherFile = null; - - try { - byte[] key = Util.getSecretBytes(64); - byte[] plaintextInput = Util.getSecretBytes(MEBIBYTE); - EncryptResult encryptResult = encryptData(plaintextInput, key, incremental); - - cipherFile = writeToFile(encryptResult.ciphertext); - - AttachmentCipherInputStream.createForAttachment(cipherFile, plaintextInput.length, key, null, encryptResult.incrementalDigest, encryptResult.chunkSizeChoice); - } finally { - if (cipherFile != null) { - cipherFile.delete(); - } - } - } - - @Test(expected = InvalidMessageException.class) - public void attachment_decryptFailOnBadDigest_nonIncremental() throws IOException, InvalidMessageException { - attachment_decryptFailOnBadDigest(false); - } - - @Test(expected = InvalidMessageException.class) - public void attachment_decryptFailOnBadDigest_incremental() throws IOException, InvalidMessageException { - attachment_decryptFailOnBadDigest(true); - } - - private void attachment_decryptFailOnBadDigest(boolean incremental) throws IOException, InvalidMessageException { - File cipherFile = null; - - try { - byte[] key = Util.getSecretBytes(64); - byte[] plaintextInput = Util.getSecretBytes(MEBIBYTE); - EncryptResult encryptResult = encryptData(plaintextInput, key, incremental); - byte[] badDigest = new byte[32]; - - cipherFile = writeToFile(encryptResult.ciphertext); - - InputStream stream = AttachmentCipherInputStream.createForAttachment(cipherFile, plaintextInput.length, key, badDigest, encryptResult.incrementalDigest, encryptResult.chunkSizeChoice); - - // In incremental mode, we'll only check the digest after reading the whole thing - if (incremental) { - StreamUtil.readFully(stream); - } - } finally { - if (cipherFile != null) { - cipherFile.delete(); - } - } - } - - @Test - public void attachment_decryptFailOnBadIncrementalDigest() throws IOException { - File cipherFile = null; - boolean hitCorrectException = false; - - try { - byte[] key = Util.getSecretBytes(64); - byte[] plaintextInput = Util.getSecretBytes(MEBIBYTE); - - EncryptResult encryptResult = encryptData(plaintextInput, key, true); - byte[] badDigest = Util.getSecretBytes(encryptResult.incrementalDigest.length); - - cipherFile = writeToFile(encryptResult.ciphertext); - - - InputStream decryptedStream = AttachmentCipherInputStream.createForAttachment(cipherFile, plaintextInput.length, key, encryptResult.digest, badDigest, encryptResult.chunkSizeChoice); - byte[] plaintextOutput = readInputStreamFully(decryptedStream); - fail(); - } catch (InvalidMacException e) { - hitCorrectException = true; - } catch (InvalidMessageException e) { - hitCorrectException = false; - } finally { - if (cipherFile != null) { - cipherFile.delete(); - } - } - - assertTrue(hitCorrectException); - } - - @Test - public void attachment_encryptDecryptPaddedContent() throws IOException, InvalidMessageException { - int[] lengths = { 531, 600, 724, 1019, 1024 }; - - for (int length : lengths) { - byte[] plaintextInput = new byte[length]; - - for (int i = 0; i < length; i++) { - plaintextInput[i] = (byte) 0x97; - } - - byte[] key = Util.getSecretBytes(64); - byte[] iv = Util.getSecretBytes(16); - ByteArrayInputStream inputStream = new ByteArrayInputStream(plaintextInput); - InputStream paddedInputStream = new PaddingInputStream(inputStream, length); - ByteArrayOutputStream destinationOutputStream = new ByteArrayOutputStream(); - - DigestingOutputStream encryptingOutputStream = new AttachmentCipherOutputStreamFactory(key, iv).createFor(destinationOutputStream); - - Util.copy(paddedInputStream, encryptingOutputStream); - - encryptingOutputStream.flush(); - encryptingOutputStream.close(); - - byte[] encryptedData = destinationOutputStream.toByteArray(); - byte[] digest = encryptingOutputStream.getTransmittedDigest(); - - File cipherFile = writeToFile(encryptedData); - - InputStream decryptedStream = AttachmentCipherInputStream.createForAttachment(cipherFile, length, key, digest, null, 0); - byte[] plaintextOutput = readInputStreamFully(decryptedStream); - - assertArrayEquals(plaintextInput, plaintextOutput); - - cipherFile.delete(); - } - } - - @Test - public void archive_encryptDecrypt() throws IOException, InvalidMessageException { - byte[] key = Util.getSecretBytes(64); - MediaRootBackupKey.MediaKeyMaterial keyMaterial = AttachmentCipherTestHelper.createMediaKeyMaterial(key); - byte[] plaintextInput = "Peter Parker".getBytes(); - EncryptResult encryptResult = encryptData(plaintextInput, key, false); - File cipherFile = writeToFile(encryptResult.ciphertext); - InputStream inputStream = AttachmentCipherInputStream.createForArchivedMedia(keyMaterial, cipherFile, plaintextInput.length); - byte[] plaintextOutput = readInputStreamFully(inputStream); - - assertArrayEquals(plaintextInput, plaintextOutput); - - cipherFile.delete(); - } - - @Test - public void archive_encryptDecryptEmpty() throws IOException, InvalidMessageException { - byte[] key = Util.getSecretBytes(64); - MediaRootBackupKey.MediaKeyMaterial keyMaterial = AttachmentCipherTestHelper.createMediaKeyMaterial(key); - byte[] plaintextInput = "".getBytes(); - EncryptResult encryptResult = encryptData(plaintextInput, key, false); - File cipherFile = writeToFile(encryptResult.ciphertext); - InputStream inputStream = AttachmentCipherInputStream.createForArchivedMedia(keyMaterial, cipherFile, plaintextInput.length); - byte[] plaintextOutput = readInputStreamFully(inputStream); - - assertArrayEquals(plaintextInput, plaintextOutput); - - cipherFile.delete(); - } - - @Test - public void archive_decryptFailOnBadKey() throws IOException { - File cipherFile = null; - boolean hitCorrectException = false; - - try { - byte[] key = Util.getSecretBytes(64); - byte[] badKey = Util.getSecretBytes(64); - MediaRootBackupKey.MediaKeyMaterial keyMaterial = AttachmentCipherTestHelper.createMediaKeyMaterial(badKey); - byte[] plaintextInput = "Gwen Stacy".getBytes(); - EncryptResult encryptResult = encryptData(plaintextInput, key, false); - - cipherFile = writeToFile(encryptResult.ciphertext); - - AttachmentCipherInputStream.createForArchivedMedia(keyMaterial, cipherFile, plaintextInput.length); - } catch (InvalidMessageException e) { - hitCorrectException = true; - } finally { - if (cipherFile != null) { - cipherFile.delete(); - } - } - - assertTrue(hitCorrectException); - } - - @Test - public void archive_encryptDecryptPaddedContent() throws IOException, InvalidMessageException { - int[] lengths = { 531, 600, 724, 1019, 1024 }; - - for (int length : lengths) { - byte[] plaintextInput = new byte[length]; - - for (int i = 0; i < length; i++) { - plaintextInput[i] = (byte) 0x97; - } - - byte[] key = Util.getSecretBytes(64); - byte[] iv = Util.getSecretBytes(16); - ByteArrayInputStream inputStream = new ByteArrayInputStream(plaintextInput); - InputStream paddedInputStream = new PaddingInputStream(inputStream, length); - ByteArrayOutputStream destinationOutputStream = new ByteArrayOutputStream(); - - DigestingOutputStream encryptingOutputStream = new AttachmentCipherOutputStreamFactory(key, iv).createFor(destinationOutputStream); - - Util.copy(paddedInputStream, encryptingOutputStream); - - encryptingOutputStream.flush(); - encryptingOutputStream.close(); - - byte[] encryptedData = destinationOutputStream.toByteArray(); - - File cipherFile = writeToFile(encryptedData); - - MediaRootBackupKey.MediaKeyMaterial keyMaterial = AttachmentCipherTestHelper.createMediaKeyMaterial(key); - InputStream decryptedStream = AttachmentCipherInputStream.createForArchivedMedia(keyMaterial, cipherFile, length); - byte[] plaintextOutput = readInputStreamFully(decryptedStream); - - assertArrayEquals(plaintextInput, plaintextOutput); - - cipherFile.delete(); - } - } - - @Test - public void archive_decryptFailOnBadMac() throws IOException { - File cipherFile = null; - boolean hitCorrectException = false; - - try { - byte[] key = Util.getSecretBytes(64); - byte[] plaintextInput = Util.getSecretBytes(MEBIBYTE); - EncryptResult encryptResult = encryptData(plaintextInput, key, true); - byte[] badMacCiphertext = Arrays.copyOf(encryptResult.ciphertext, encryptResult.ciphertext.length); - - badMacCiphertext[badMacCiphertext.length - 1] += 1; - - cipherFile = writeToFile(badMacCiphertext); - - MediaRootBackupKey.MediaKeyMaterial keyMaterial = AttachmentCipherTestHelper.createMediaKeyMaterial(key); - AttachmentCipherInputStream.createForArchivedMedia(keyMaterial, cipherFile, plaintextInput.length); - fail(); - } catch (InvalidMessageException e) { - hitCorrectException = true; - } finally { - if (cipherFile != null) { - cipherFile.delete(); - } - } - - assertTrue(hitCorrectException); - } - - @Test - public void sticker_encryptDecrypt() throws IOException, InvalidMessageException { - assumeLibSignalSupportedOnOS(); - - byte[] packKey = Util.getSecretBytes(32); - byte[] plaintextInput = Util.getSecretBytes(MEBIBYTE); - EncryptResult encryptResult = encryptData(plaintextInput, expandPackKey(packKey), true); - InputStream inputStream = AttachmentCipherInputStream.createForStickerData(encryptResult.ciphertext, packKey); - byte[] plaintextOutput = readInputStreamFully(inputStream); - - assertArrayEquals(plaintextInput, plaintextOutput); - } - - @Test - public void sticker_encryptDecryptEmpty() throws IOException, InvalidMessageException { - assumeLibSignalSupportedOnOS(); - - byte[] packKey = Util.getSecretBytes(32); - byte[] plaintextInput = "".getBytes(); - EncryptResult encryptResult = encryptData(plaintextInput, expandPackKey(packKey), true); - InputStream inputStream = AttachmentCipherInputStream.createForStickerData(encryptResult.ciphertext, packKey); - byte[] plaintextOutput = readInputStreamFully(inputStream); - - assertArrayEquals(plaintextInput, plaintextOutput); - } - - @Test - public void sticker_decryptFailOnBadKey() throws IOException { - assumeLibSignalSupportedOnOS(); - - boolean hitCorrectException = false; - - try { - byte[] packKey = Util.getSecretBytes(32); - byte[] plaintextInput = Util.getSecretBytes(MEBIBYTE); - EncryptResult encryptResult = encryptData(plaintextInput, expandPackKey(packKey), true); - byte[] badPackKey = new byte[32]; - - AttachmentCipherInputStream.createForStickerData(encryptResult.ciphertext, badPackKey); - } catch (InvalidMessageException e) { - hitCorrectException = true; - } - - assertTrue(hitCorrectException); - } - - @Test - public void sticker_decryptFailOnBadMac() throws IOException { - assumeLibSignalSupportedOnOS(); - - boolean hitCorrectException = false; - - try { - byte[] packKey = Util.getSecretBytes(32); - byte[] plaintextInput = Util.getSecretBytes(MEBIBYTE); - EncryptResult encryptResult = encryptData(plaintextInput, expandPackKey(packKey), true); - byte[] badMacCiphertext = Arrays.copyOf(encryptResult.ciphertext, encryptResult.ciphertext.length); - - badMacCiphertext[badMacCiphertext.length - 1] += 1; - - AttachmentCipherInputStream.createForStickerData(badMacCiphertext, packKey); - } catch (InvalidMessageException e) { - hitCorrectException = true; - } - - assertTrue(hitCorrectException); - } - - private static EncryptResult encryptData(byte[] data, byte[] keyMaterial, boolean withIncremental) throws IOException { - ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); - ByteArrayOutputStream incrementalDigestOut = new ByteArrayOutputStream(); - byte[] iv = Util.getSecretBytes(16); - AttachmentCipherOutputStreamFactory factory = new AttachmentCipherOutputStreamFactory(keyMaterial, iv); - - DigestingOutputStream encryptStream; - final ChunkSizeChoice sizeChoice = ChunkSizeChoice.inferChunkSize(data.length); - if (withIncremental) { - encryptStream = factory.createIncrementalFor(outputStream, data.length, sizeChoice, incrementalDigestOut); - } else { - encryptStream = factory.createFor(outputStream); - } - - encryptStream.write(data); - encryptStream.flush(); - encryptStream.close(); - incrementalDigestOut.close(); - - return new EncryptResult(outputStream.toByteArray(), encryptStream.getTransmittedDigest(), incrementalDigestOut.toByteArray(), sizeChoice.getSizeInBytes()); - } - - private static File writeToFile(byte[] data) throws IOException { - File file = File.createTempFile("temp", ".data"); - OutputStream outputStream = new FileOutputStream(file); - - outputStream.write(data); - outputStream.close(); - - return file; - } - - private static byte[] readInputStreamFully(InputStream inputStream) throws IOException { - return Util.readFullyAsBytes(inputStream); - } - - private static byte[] expandPackKey(byte[] shortKey) { - return HKDF.deriveSecrets(shortKey, "Sticker Pack".getBytes(), 64); - } - - private static class EncryptResult { - final byte[] ciphertext; - final byte[] digest; - final byte[] incrementalDigest; - final int chunkSizeChoice; - - private EncryptResult(byte[] ciphertext, byte[] digest, byte[] incrementalDigest, int chunkSizeChoice) { - this.ciphertext = ciphertext; - this.digest = digest; - this.incrementalDigest = incrementalDigest; - this.chunkSizeChoice = chunkSizeChoice; - } - } -} diff --git a/libsignal-service/src/test/java/org/whispersystems/signalservice/api/crypto/AttachmentCipherTest.kt b/libsignal-service/src/test/java/org/whispersystems/signalservice/api/crypto/AttachmentCipherTest.kt new file mode 100644 index 0000000000..0d9e2f7140 --- /dev/null +++ b/libsignal-service/src/test/java/org/whispersystems/signalservice/api/crypto/AttachmentCipherTest.kt @@ -0,0 +1,548 @@ +package org.whispersystems.signalservice.api.crypto + +import assertk.assertThat +import assertk.assertions.isEqualTo +import assertk.assertions.isTrue +import assertk.fail +import org.conscrypt.Conscrypt +import org.junit.Assert +import org.junit.Test +import org.signal.core.util.StreamUtil +import org.signal.core.util.copyTo +import org.signal.libsignal.protocol.InvalidMessageException +import org.signal.libsignal.protocol.incrementalmac.ChunkSizeChoice +import org.signal.libsignal.protocol.incrementalmac.InvalidMacException +import org.signal.libsignal.protocol.kdf.HKDF +import org.whispersystems.signalservice.api.crypto.AttachmentCipherTestHelper.createMediaKeyMaterial +import org.whispersystems.signalservice.internal.crypto.PaddingInputStream +import org.whispersystems.signalservice.internal.push.http.AttachmentCipherOutputStreamFactory +import org.whispersystems.signalservice.internal.util.Util +import org.whispersystems.signalservice.testutil.LibSignalLibraryUtil +import java.io.ByteArrayInputStream +import java.io.ByteArrayOutputStream +import java.io.File +import java.io.FileOutputStream +import java.io.IOException +import java.io.InputStream +import java.io.OutputStream +import java.lang.AssertionError +import java.security.Security +import java.util.Random + +class AttachmentCipherTest { + @Test + @Throws(IOException::class, InvalidMessageException::class) + fun attachment_encryptDecrypt_nonIncremental() { + attachment_encryptDecrypt(incremental = false, fileSize = MEBIBYTE) + } + + @Test + @Throws(IOException::class, InvalidMessageException::class) + fun attachment_encryptDecrypt_incremental() { + attachment_encryptDecrypt(incremental = true, fileSize = MEBIBYTE) + } + + @Test + @Throws(IOException::class, InvalidMessageException::class) + fun attachment_encryptDecrypt_incremental_manyFileSizes() { + // Designed to stress the various boundary conditions of reading the final mac + for (i in 0..99) { + attachment_encryptDecrypt(incremental = true, fileSize = MEBIBYTE + Random().nextInt(1, 64 * 1024)) + } + } + + @Throws(IOException::class, InvalidMessageException::class) + private fun attachment_encryptDecrypt(incremental: Boolean, fileSize: Int) { + val key = Util.getSecretBytes(64) + val plaintextInput = Util.getSecretBytes(fileSize) + + val encryptResult = encryptData(plaintextInput, key, incremental) + val cipherFile = writeToFile(encryptResult.ciphertext) + + val inputStream: InputStream = AttachmentCipherInputStream.createForAttachment(cipherFile, plaintextInput.size.toLong(), key, encryptResult.digest, encryptResult.incrementalDigest, encryptResult.chunkSizeChoice) + val plaintextOutput = readInputStreamFully(inputStream) + + assertThat(plaintextOutput).isEqualTo(plaintextInput) + + cipherFile.delete() + } + + @Test + @Throws(IOException::class, InvalidMessageException::class) + fun attachment_encryptDecryptEmpty_nonIncremental() { + attachment_encryptDecryptEmpty(incremental = false) + } + + @Test + @Throws(IOException::class, InvalidMessageException::class) + fun attachment_encryptDecryptEmpty_incremental() { + attachment_encryptDecryptEmpty(incremental = true) + } + + @Throws(IOException::class, InvalidMessageException::class) + private fun attachment_encryptDecryptEmpty(incremental: Boolean) { + val key = Util.getSecretBytes(64) + val plaintextInput = "".toByteArray() + + val encryptResult = encryptData(plaintextInput, key, incremental) + val cipherFile = writeToFile(encryptResult.ciphertext) + + val inputStream: InputStream = AttachmentCipherInputStream.createForAttachment(cipherFile, plaintextInput.size.toLong(), key, encryptResult.digest, encryptResult.incrementalDigest, encryptResult.chunkSizeChoice) + val plaintextOutput = readInputStreamFully(inputStream) + + Assert.assertArrayEquals(plaintextInput, plaintextOutput) + + cipherFile.delete() + } + + @Test(expected = InvalidMessageException::class) + @Throws(IOException::class, InvalidMessageException::class) + fun attachment_decryptFailOnBadKey_nonIncremental() { + attachment_decryptFailOnBadKey(incremental = false) + } + + @Test(expected = InvalidMessageException::class) + @Throws(IOException::class, InvalidMessageException::class) + fun attachment_decryptFailOnBadKey_incremental() { + attachment_decryptFailOnBadKey(incremental = true) + } + + @Throws(IOException::class, InvalidMessageException::class) + private fun attachment_decryptFailOnBadKey(incremental: Boolean) { + var cipherFile: File? = null + + try { + val key = Util.getSecretBytes(64) + val plaintextInput = Util.getSecretBytes(MEBIBYTE) + + val encryptResult = encryptData(plaintextInput, key, incremental) + cipherFile = writeToFile(encryptResult.ciphertext) + + val badKey = ByteArray(64) + AttachmentCipherInputStream.createForAttachment(cipherFile, plaintextInput.size.toLong(), badKey, encryptResult.digest, null, 0) + } finally { + cipherFile?.delete() + } + } + + @Test(expected = InvalidMessageException::class) + @Throws(IOException::class, InvalidMessageException::class) + fun attachment_decryptFailOnBadMac_nonIncremental() { + attachment_decryptFailOnBadMac(incremental = false) + } + + @Test(expected = InvalidMessageException::class) + @Throws(IOException::class, InvalidMessageException::class) + fun attachment_decryptFailOnBadMac_incremental() { + attachment_decryptFailOnBadMac(incremental = true) + } + + @Throws(IOException::class, InvalidMessageException::class) + private fun attachment_decryptFailOnBadMac(incremental: Boolean) { + var cipherFile: File? = null + + try { + val key = Util.getSecretBytes(64) + val plaintextInput = Util.getSecretBytes(MEBIBYTE) + + val encryptResult = encryptData(plaintextInput, key, incremental) + val badMacCiphertext = encryptResult.ciphertext.copyOf(encryptResult.ciphertext.size) + + badMacCiphertext[badMacCiphertext.size - 1] = (badMacCiphertext[badMacCiphertext.size - 1] + 1).toByte() + + cipherFile = writeToFile(badMacCiphertext) + + val stream: InputStream = AttachmentCipherInputStream.createForAttachment(cipherFile, plaintextInput.size.toLong(), key, encryptResult.digest, encryptResult.incrementalDigest, encryptResult.chunkSizeChoice) + + // In incremental mode, we'll only check the digest after reading the whole thing + if (incremental) { + StreamUtil.readFully(stream) + } + } finally { + cipherFile?.delete() + } + } + + @Test(expected = InvalidMessageException::class) + @Throws(IOException::class, InvalidMessageException::class) + fun attachment_decryptFailOnNullDigest_nonIncremental() { + attachment_decryptFailOnNullDigest(incremental = false) + } + + @Test(expected = InvalidMessageException::class) + @Throws(IOException::class, InvalidMessageException::class) + fun attachment_decryptFailOnNullDigest_incremental() { + attachment_decryptFailOnNullDigest(incremental = true) + } + + @Throws(IOException::class, InvalidMessageException::class) + private fun attachment_decryptFailOnNullDigest(incremental: Boolean) { + var cipherFile: File? = null + + try { + val key = Util.getSecretBytes(64) + val plaintextInput = Util.getSecretBytes(MEBIBYTE) + val encryptResult = encryptData(plaintextInput, key, incremental) + + cipherFile = writeToFile(encryptResult.ciphertext) + + AttachmentCipherInputStream.createForAttachment(cipherFile, plaintextInput.size.toLong(), key, null, encryptResult.incrementalDigest, encryptResult.chunkSizeChoice) + } finally { + cipherFile?.delete() + } + } + + @Test(expected = InvalidMessageException::class) + @Throws(IOException::class, InvalidMessageException::class) + fun attachment_decryptFailOnBadDigest_nonIncremental() { + attachment_decryptFailOnBadDigest(incremental = false) + } + + @Test(expected = InvalidMessageException::class) + @Throws(IOException::class, InvalidMessageException::class) + fun attachment_decryptFailOnBadDigest_incremental() { + attachment_decryptFailOnBadDigest(incremental = true) + } + + @Throws(IOException::class, InvalidMessageException::class) + private fun attachment_decryptFailOnBadDigest(incremental: Boolean) { + var cipherFile: File? = null + + try { + val key = Util.getSecretBytes(64) + val plaintextInput = Util.getSecretBytes(MEBIBYTE) + + val encryptResult = encryptData(plaintextInput, key, incremental) + val badDigest = ByteArray(32) + + cipherFile = writeToFile(encryptResult.ciphertext) + + val stream: InputStream = AttachmentCipherInputStream.createForAttachment(cipherFile, plaintextInput.size.toLong(), key, badDigest, encryptResult.incrementalDigest, encryptResult.chunkSizeChoice) + + // In incremental mode, we'll only check the digest after reading the whole thing + if (incremental) { + StreamUtil.readFully(stream) + } + } finally { + cipherFile?.delete() + } + } + + @Test + @Throws(IOException::class) + fun attachment_decryptFailOnBadIncrementalDigest() { + var cipherFile: File? = null + var hitCorrectException = false + + try { + val key = Util.getSecretBytes(64) + val plaintextInput = Util.getSecretBytes(MEBIBYTE) + + val encryptResult = encryptData(plaintextInput, key, true) + val badDigest = Util.getSecretBytes(encryptResult.incrementalDigest.size) + + cipherFile = writeToFile(encryptResult.ciphertext) + + val decryptedStream: InputStream = AttachmentCipherInputStream.createForAttachment(cipherFile, plaintextInput.size.toLong(), key, encryptResult.digest, badDigest, encryptResult.chunkSizeChoice) + val plaintextOutput = readInputStreamFully(decryptedStream) + + fail(AssertionError("Expected to fail before hitting this line")) + } catch (e: InvalidMacException) { + hitCorrectException = true + } catch (e: InvalidMessageException) { + hitCorrectException = false + } finally { + cipherFile?.delete() + } + + assertThat(hitCorrectException).isTrue() + } + + @Test + @Throws(IOException::class, InvalidMessageException::class) + fun attachment_encryptDecryptPaddedContent() { + val lengths = intArrayOf(531, 600, 724, 1019, 1024) + + for (length in lengths) { + val plaintextInput = ByteArray(length) + + for (i in 0..