chacha: Reduce API surface for Counter/IV management.

Get rid of the `encrypt_in_place` wrapper.

Replace most of the IV/Counter API with a simpler, private one.

This makes it clearer that the callers are handling nonces correctly.
This will also make it easier to refactor the Counter handling later.

src/aead/chacha.rs

@@ -49,22 +49,29 @@ impl Key {
         self.encrypt_within(counter, Overlapping::in_place(in_out))
     }
 
+    // Encrypts `in_out` with the counter 0 and returns counter 1,
+    // where the counter is derived from the nonce `nonce`.
     #[inline]
-    pub fn encrypt_iv_xor_in_place(&self, iv: Iv, in_out: &mut [u8; 32]) {
-        // It is safe to use `into_counter_for_single_block_less_safe()`
-        // because `in_out` is exactly one block long.
-        debug_assert!(in_out.len() <= BLOCK_LEN);
-        self.encrypt_in_place(iv.into_counter_for_single_block_less_safe(), in_out);
+    pub fn encrypt_single_block_with_ctr_0<const N: usize>(
+        &self,
+        nonce: Nonce,
+        in_out: &mut [u8; N],
+    ) -> Counter {
+        assert!(N <= BLOCK_LEN);
+        let (zero, one) = Counter::zero_one_less_safe(nonce);
+        self.encrypt_within(zero, Overlapping::in_place(in_out));
+        one
     }
 
     #[inline]
     pub fn new_mask(&self, sample: Sample) -> [u8; 5] {
-        let mut out: [u8; 5] = [0; 5];
-        let iv = Iv::assume_unique_for_key(sample);
-
-        debug_assert!(out.len() <= BLOCK_LEN);
-        self.encrypt_in_place(iv.into_counter_for_single_block_less_safe(), &mut out);
+        let (ctr, nonce) = sample.split_at(4);
+        let ctr = u32::from_le_bytes(ctr.try_into().unwrap());
+        let nonce = Nonce::assume_unique_for_key(nonce.try_into().unwrap());
+        let ctr = Counter::from_nonce_and_ctr(nonce, ctr);
 
+        let mut out: [u8; 5] = [0; 5];
+        self.encrypt_within(ctr, Overlapping::in_place(&mut out));
         out
     }
 
@@ -128,21 +135,19 @@ impl Key {
 pub struct Counter([u32; 4]);
 
 impl Counter {
-    pub fn zero(nonce: Nonce) -> Self {
-        Self::from_nonce_and_ctr(nonce, 0)
+    // Nonce-reuse: the caller must only use the first counter (0) for at most
+    // a single block.
+    fn zero_one_less_safe(nonce: Nonce) -> (Self, Self) {
+        let ctr0 @ Self([_, n0, n1, n2]) = Self::from_nonce_and_ctr(nonce, 0);
+        let ctr1 = Self([1, n0, n1, n2]);
+        (ctr0, ctr1)
     }
 
     fn from_nonce_and_ctr(nonce: Nonce, ctr: u32) -> Self {
         let [n0, n1, n2] = nonce.as_ref().array_split_map(u32::from_le_bytes);
         Self([ctr, n0, n1, n2])
     }
 
-    pub fn increment(&mut self) -> Iv {
-        let iv = Iv(self.0);
-        self.0[0] += 1;
-        iv
-    }
-
     /// This is "less safe" because it hands off management of the counter to
     /// the caller.
     #[cfg(any(
@@ -159,21 +164,6 @@ impl Counter {
     }
 }
 
-/// The IV for a single block encryption.
-///
-/// Intentionally not `Clone` to ensure each is used only once.
-pub struct Iv([u32; 4]);
-
-impl Iv {
-    fn assume_unique_for_key(value: [u8; 16]) -> Self {
-        Self(value.array_split_map(u32::from_le_bytes))
-    }
-
-    fn into_counter_for_single_block_less_safe(self) -> Counter {
-        Counter(self.0)
-    }
-}
-
 pub const KEY_LEN: usize = 32;
 
 const BLOCK_LEN: usize = 64;

src/aead/chacha20_poly1305.rs

@@ -224,10 +224,8 @@ fn has_integrated(cpu_features: cpu::Features) -> bool {
 
 // Also used by chacha20_poly1305_openssh.
 pub(super) fn begin(key: &chacha::Key, nonce: Nonce) -> (Counter, poly1305::Key) {
-    let mut counter = Counter::zero(nonce);
-    let iv = counter.increment();
     let mut key_bytes = [0u8; poly1305::KEY_LEN];
-    key.encrypt_iv_xor_in_place(iv, &mut key_bytes);
+    let counter = key.encrypt_single_block_with_ctr_0(nonce, &mut key_bytes);
     let poly1305_key = poly1305::Key::new(key_bytes);
     (counter, poly1305_key)
 }

src/aead/chacha20_poly1305_openssh.rs

@@ -33,7 +33,7 @@ use super::{
     chacha::{self, *},
     chacha20_poly1305, cpu, poly1305, Nonce, Tag,
 };
-use crate::{constant_time, error};
+use crate::{constant_time, error, polyfill::slice};
 
 /// A key for sealing packets.
 pub struct SealingKey {
@@ -65,16 +65,18 @@ impl SealingKey {
         plaintext_in_ciphertext_out: &mut [u8],
         tag_out: &mut [u8; TAG_LEN],
     ) {
-        let (len_in_out, data_and_padding_in_out) =
-            plaintext_in_ciphertext_out.split_at_mut(PACKET_LENGTH_LEN);
+        // XXX/TODO(SemVer): Refactor API to return an error.
+        let (len_in_out, data_and_padding_in_out): (&mut [u8; PACKET_LENGTH_LEN], _) =
+            slice::split_first_chunk_mut(plaintext_in_ciphertext_out).unwrap();
 
         let cpu_features = cpu::features();
         let (counter, poly_key) =
             chacha20_poly1305::begin(&self.key.k_2, make_nonce(sequence_number));
 
-        self.key
+        let _: Counter = self
+            .key
             .k_1
-            .encrypt_in_place(make_counter(sequence_number), len_in_out);
+            .encrypt_single_block_with_ctr_0(make_nonce(sequence_number), len_in_out);
         self.key
             .k_2
             .encrypt_in_place(counter, data_and_padding_in_out);
@@ -107,8 +109,10 @@ impl OpeningKey {
         encrypted_packet_length: [u8; PACKET_LENGTH_LEN],
     ) -> [u8; PACKET_LENGTH_LEN] {
         let mut packet_length = encrypted_packet_length;
-        let counter = make_counter(sequence_number);
-        self.key.k_1.encrypt_in_place(counter, &mut packet_length);
+        let _: Counter = self
+            .key
+            .k_1
+            .encrypt_single_block_with_ctr_0(make_nonce(sequence_number), &mut packet_length);
         packet_length
     }
 
@@ -172,11 +176,6 @@ fn make_nonce(sequence_number: u32) -> Nonce {
     Nonce::assume_unique_for_key(nonce)
 }
 
-fn make_counter(sequence_number: u32) -> Counter {
-    let nonce = make_nonce(sequence_number);
-    Counter::zero(nonce)
-}
-
 /// The length of key.
 pub const KEY_LEN: usize = chacha::KEY_LEN * 2;
 

src/polyfill/slice.rs

@@ -112,3 +112,16 @@ pub fn split_at_checked<T>(slice: &[T], i: usize) -> Option<(&[T], &[T])> {
         None
     }
 }
+
+// TODO(MSRV-1.77): Use `slice::split_first_chunk_mut`.
+#[inline(always)]
+pub fn split_first_chunk_mut<T, const N: usize>(
+    slice: &mut [T],
+) -> Option<(&mut [T; N], &mut [T])> {
+    if slice.len() >= N {
+        let (head, tail) = slice.split_at_mut(N);
+        head.try_into().ok().map(|head| (head, tail))
+    } else {
+        None
+    }
+}