From 8e4a6ebcb1f5c033c572a0b7c1326418cccd0258 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Sun, 27 Aug 2023 00:43:57 +0800
Subject: [PATCH] feat: update tests

---
 src/lib.rs | 86 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 86 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index 45c8790..c2eb352 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -11,6 +11,83 @@ mod util;
 mod encryptor;
 mod decryptor;
 
+#[test]
+fn test128() {
+    use aes_gcm::{aead::{Aead, Nonce, Payload}, Aes128Gcm, KeyInit};
+    let knp = vec![
+        ([0; 16], [0; 12], &[] as &[u8], b"Hello World!" as &[u8]),
+        ([0; 16], [0; 12], &[1; 16], b"Hello World!" as &[u8]),
+        ([0; 16], [0; 12], &[1; 17], b"Hello World!" as &[u8]),
+        ([0; 16], [0; 12], &[1; 32], b"Hello World!" as &[u8]),
+        ([0; 16], [0; 12], &[1; 64], b"Hello World!" as &[u8]),
+        ([0; 16], [0; 12], &[1, 2, 3], b"Hello World!" as &[u8]),
+        ([1; 16], [0; 12], &[], b"Hello World!"),
+        ([0; 16], [1; 12], &[], b"Hello World!"),
+        ([1; 16], [1; 12], &[], b"Hello World ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~!"),
+        ([1; 16], [1; 12], &[0; 129], b"Hello World ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~!"),
+        ([0xff; 16], [0; 12], &[], b"Hello World!"),
+        ([0; 16], [0xff; 12], &[], b"Hello World!"),
+        ([0xff; 16], [0xff; 12], &[], b"Hello World ~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~!"),
+        ([0xff; 16], [0xff; 12], &[11, 22, 33], b"Hello World ~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~!"),
+    ];
+
+    for (key, nonce, aad, plaintext) in knp {
+        // encrypt
+        let mut ciphertext = vec![];
+        let mut encryptor = Aes128GcmStreamEncryptor::new(key.clone(), &nonce);
+        if !aad.is_empty() {
+            encryptor.init_adata(aad);
+        }
+        ciphertext.extend_from_slice(&encryptor.next(plaintext));
+        let (last_block, tag) = encryptor.finalize();
+        ciphertext.extend_from_slice(&last_block);
+        ciphertext.extend_from_slice(&tag);
+
+        // decrypt 1
+        let mut decryptor = Aes128GcmStreamDecryptor::new(key.clone(), &nonce);
+        if !aad.is_empty() {
+            decryptor.init_adata(aad);
+        }
+        let mut plaintext1 = decryptor.next(ciphertext.as_slice());
+        let plaintext2 = decryptor.finalize().expect("decryptor decrypt");
+        plaintext1.extend_from_slice(&plaintext2);
+        assert_eq!(plaintext, plaintext1.as_slice());
+
+        // decrypt 2
+        let cipher = Aes128Gcm::new_from_slice(&key).expect("new from key slice");
+        let mut decrypt_nonce = Nonce::<Aes128Gcm>::default();
+        let m: &mut [u8] = decrypt_nonce.as_mut();
+        for i in 0..m.len() {
+            m[i] = nonce[i];
+        }
+        let decrypted_plaintext = if aad.is_empty() {
+            cipher.decrypt(&decrypt_nonce, ciphertext.as_slice()).expect("decrypt1")
+        } else {
+            cipher.decrypt(&decrypt_nonce, Payload {
+                msg: ciphertext.as_slice(),
+                aad,
+            }).expect("decrypt2")
+        };
+        assert_eq!(plaintext, decrypted_plaintext.as_slice());
+    }
+}
+
 #[test]
 fn test256() {
     use aes_gcm::{aead::{Aead, Nonce, Payload}, Aes256Gcm, KeyInit};
@@ -31,11 +108,20 @@ fn test256() {
         ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
         ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
         ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
         ~~~~~~~~~!"),
         ([0xff; 32], [0xff; 12], &[11, 22, 33], b"Hello World ~~~~~~~~~~~~~~~~~\
         ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
         ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
         ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
+        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\
         ~~~~~~~~~!"),
     ];