feat: Add AEAD encryption function with explicit sequence number (#40)

* Add AEAD encryption function with explicit sequence number

* Fixes

* formatting

* Fixes

Author: Not-Nik

src/aead.rs

@@ -512,6 +512,15 @@ impl Aead {
         })
     }
 
+    fn make_nonce(nonce: &mut [u8; NONCE_LEN], seq: SequenceNumber) {
+        for (n, &s) in nonce[NONCE_LEN - COUNTER_LEN..]
+            .iter_mut()
+            .zip(&seq.to_be_bytes())
+        {
+            *n ^= s;
+        }
+    }
+
     pub fn import_key(algorithm: AeadAlgorithms, key: &[u8]) -> Result<SymKey, Error> {
         let slot = p11::Slot::internal().map_err(|_| Error::Internal)?;
 
@@ -590,6 +599,49 @@ impl Aead {
         Ok(ct)
     }
 
+    /// Encrypt with an explicit sequence number. Mirrors `decrypt`'s nonce
+    /// construction: the final nonce is `nonce_base XOR encode_be(seq)` over
+    /// the trailing 8 bytes. The NSS PKCS#11 context's internal counter is
+    /// not used (`CKG_NO_GENERATE`). The caller must never reuse
+    /// `(nonce_base, seq)` with the same key.
+    pub fn encrypt_with_seq(
+        &mut self,
+        aad: &[u8],
+        seq: SequenceNumber,
+        pt: &[u8],
+    ) -> Result<Vec<u8>, Error> {
+        crate::init()?;
+
+        assert_eq!(self.mode, Mode::Encrypt);
+        let mut nonce = self.nonce_base;
+        Self::make_nonce(&mut nonce, seq);
+        let mut ct = vec![0; pt.len() + TAG_LEN];
+        let mut ct_len: c_int = 0;
+        let mut tag = vec![0; TAG_LEN];
+        secstatus_to_res(unsafe {
+            PK11_AEADOp(
+                *self.ctx,
+                CK_GENERATOR_FUNCTION::from(CKG_NO_GENERATE),
+                c_int_len(NONCE_LEN - COUNTER_LEN)?,
+                nonce.as_mut_ptr(),
+                c_int_len(nonce.len())?,
+                aad.as_ptr(),
+                c_int_len(aad.len())?,
+                ct.as_mut_ptr(),
+                &raw mut ct_len,
+                c_int_len(ct.len())?,
+                tag.as_mut_ptr(),
+                c_int_len(tag.len())?,
+                pt.as_ptr(),
+                c_int_len(pt.len())?,
+            )
+        })?;
+        ct.truncate(usize::try_from(ct_len).map_err(|_| Error::IntegerOverflow)?);
+        debug_assert_eq!(ct.len(), pt.len());
+        ct.append(&mut tag);
+        Ok(ct)
+    }
+
     pub fn decrypt(
         &mut self,
         aad: &[u8],
@@ -600,9 +652,7 @@ impl Aead {
 
         assert_eq!(self.mode, Mode::Decrypt);
         let mut nonce = self.nonce_base;
-        for (i, n) in nonce.iter_mut().rev().take(COUNTER_LEN).enumerate() {
-            *n ^= u8::try_from((seq >> (8 * i)) & 0xff).map_err(|_| Error::IntegerOverflow)?;
-        }
+        Self::make_nonce(&mut nonce, seq);
         let mut pt = vec![0; ct.len()]; // NSS needs more space than it uses for plaintext.
         let mut pt_len: c_int = 0;
         let pt_expected = ct.len().checked_sub(TAG_LEN).ok_or(Error::AeadTruncated)?;
@@ -767,4 +817,39 @@ mod test {
         // Now use the real nonce and sequence number from the example.
         decrypt(ALG, KEY, NONCE_BASE, 654_360_564, AAD, PT, CT);
     }
+
+    fn roundtrip_encrypt_with_seq(algorithm: AeadAlgorithms, key: &[u8]) {
+        const NONCE_BASE: [u8; NONCE_LEN] = [0; NONCE_LEN];
+        const AAD: &[u8] = b"associated";
+        const PT: &[u8] = b"hello sframe";
+        const SEQ: SequenceNumber = 0x0123_4567_89ab;
+
+        fixture_init();
+
+        let k = Aead::import_key(algorithm, key).unwrap();
+        let mut enc = Aead::new(Mode::Encrypt, algorithm, &k, NONCE_BASE).unwrap();
+        let ct = enc.encrypt_with_seq(AAD, SEQ, PT).unwrap();
+
+        let mut dec = Aead::new(Mode::Decrypt, algorithm, &k, NONCE_BASE).unwrap();
+        let pt = dec.decrypt(AAD, SEQ, &ct).unwrap();
+        assert_eq!(&pt[..], PT);
+    }
+
+    #[test]
+    fn encrypt_with_seq_aes128gcm() {
+        const KEY: &[u8] = &[0x42; 16];
+        roundtrip_encrypt_with_seq(AeadAlgorithms::Aes128Gcm, KEY);
+    }
+
+    #[test]
+    fn encrypt_with_seq_aes256gcm() {
+        const KEY: &[u8] = &[0x42; 32];
+        roundtrip_encrypt_with_seq(AeadAlgorithms::Aes256Gcm, KEY);
+    }
+
+    #[test]
+    fn encrypt_with_seq_chacha20poly1305() {
+        const KEY: &[u8] = &[0x42; 32];
+        roundtrip_encrypt_with_seq(AeadAlgorithms::ChaCha20Poly1305, KEY);
+    }
 }