Add varying write size test

aws · Sep 11, 2024 · 4a98ff0 · 4a98ff0
1 parent 22f6dc9
commit 4a98ff0
Showing 1 changed file with 42 additions and 3 deletions.
diff --git a/crypto/pkcs7/pkcs7_internal_bio_test.cc b/crypto/pkcs7/pkcs7_internal_bio_test.cc
@@ -25,8 +25,8 @@
 TEST(PKCS7Test, CipherBIO) {
   uint8_t key[EVP_MAX_KEY_LENGTH];
   uint8_t iv[EVP_MAX_IV_LENGTH];
-  uint8_t pt[1024 * 64 + 7];// TODO [childw] need both block-aligned, block
-                            // misaligned, small, large, etc.
+  uint8_t pt[1024 * 64 + 7];  // TODO [childw] need both block-aligned, block
+                              // misaligned, small, large, etc.
   uint8_t pt_decrypted[sizeof(pt)];
   uint8_t ct[sizeof(pt)];
   EVP_CIPHER_CTX *ctx;
@@ -43,6 +43,7 @@ TEST(PKCS7Test, CipherBIO) {
   // Unsupported CTRL flags
   EXPECT_FALSE(BIO_ctrl(bio_cipher.get(), BIO_C_DO_STATE_MACHINE, 0, NULL));
   EXPECT_FALSE(BIO_ctrl(bio_cipher.get(), BIO_CTRL_DUP, 0, NULL));
+  EXPECT_FALSE(BIO_ctrl(bio_cipher.get(), BIO_CTRL_GET_CLOSE, 0, NULL));
   // Get underlying cipher context so we can initialize it
   EXPECT_TRUE(BIO_get_cipher_ctx(bio_cipher.get(), &ctx));
   ASSERT_TRUE(
@@ -76,7 +77,7 @@ TEST(PKCS7Test, CipherBIO) {
   EXPECT_TRUE(BIO_get_cipher_status(bio_cipher.get()));
   EXPECT_EQ(Bytes(pt, sizeof(pt)), Bytes(pt_decrypted, sizeof(pt_decrypted)));
 
-  // Rounda-trip using |BIO_write| for encryption with same BIOs, reset between
+  // Round-trip using |BIO_write| for encryption with same BIOs, reset between
   // encryption/decryption using |BIO_reset|.
   bio_cipher.reset(BIO_new(BIO_f_cipher()));
   ASSERT_TRUE(bio_cipher);
@@ -111,4 +112,42 @@ TEST(PKCS7Test, CipherBIO) {
   EXPECT_TRUE(BIO_read(bio_cipher.get(), pt_decrypted, sizeof(pt_decrypted)));
   EXPECT_TRUE(BIO_get_cipher_status(bio_cipher.get()));
   EXPECT_EQ(Bytes(pt, sizeof(pt)), Bytes(pt_decrypted, sizeof(pt_decrypted)));
+
+  // TODO [childw]
+  bio_cipher.reset(BIO_new(BIO_f_cipher()));
+  ASSERT_TRUE(bio_cipher);
+  EXPECT_TRUE(BIO_get_cipher_ctx(bio_cipher.get(), &ctx));
+  ASSERT_TRUE(
+      EVP_CipherInit_ex(ctx, EVP_aes_128_gcm(), NULL, key, iv, /*enc*/ 1));
+  bio_mem.reset(BIO_new(BIO_s_mem()));
+  ASSERT_TRUE(bio_mem);
+  ASSERT_TRUE(BIO_push(bio_cipher.get(), bio_mem.get()));
+
+  std::vector<uint8_t> pt_vec, ct_vec, decrypted_pt_vec;
+  uint8_t buff[1024*1024];
+  for (size_t wsize : (size_t[]){1, 3, 7, 8, 64, 7, 0, 923, sizeof(buff), 1, 8}) {
+    ASSERT_TRUE(RAND_bytes(buff, wsize));
+    pt_vec.insert(pt_vec.end(), buff, buff + wsize);
+    EXPECT_TRUE(BIO_write(bio_cipher.get(), buff, wsize) || wsize == 0);
+  }
+  EXPECT_TRUE(BIO_flush(bio_cipher.get()));
+  EXPECT_TRUE(BIO_get_cipher_status(bio_cipher.get()));
+  while (!BIO_eof(bio_mem.get())) {
+    size_t bytes_read = BIO_read(bio_mem.get(), buff, sizeof(buff));
+    ct_vec.insert(ct_vec.end(), buff, buff + bytes_read);
+  }
+
+  EXPECT_TRUE(BIO_reset(bio_cipher.get()));  // also resets owned |bio_mem|
+  EXPECT_TRUE(BIO_write(bio_mem.get(), ct_vec.data(), ct_vec.size()));  // replace ct
+  bio_mem.release();  // |bio_cipher| took ownership
+  EXPECT_TRUE(BIO_get_cipher_ctx(bio_cipher.get(), &ctx));
+  ASSERT_TRUE(
+      EVP_CipherInit_ex(ctx, EVP_aes_128_gcm(), NULL, key, iv, /*enc*/ 0));
+  decrypted_pt_vec.resize(pt_vec.size());
+  // TODO [childw] variable read sizes?
+  EXPECT_TRUE(BIO_read(bio_cipher.get(), decrypted_pt_vec.data(), decrypted_pt_vec.size()));
+  EXPECT_TRUE(BIO_get_cipher_status(bio_cipher.get()));
+  EXPECT_EQ(pt_vec.size(), decrypted_pt_vec.size());
+  EXPECT_EQ(Bytes(pt_vec.data(), pt_vec.size()),
+            Bytes(decrypted_pt_vec.data(), decrypted_pt_vec.size()));
 }