chacha20: 64-bit counter support (#439)

Closes #334

Co-authored-by: Артём Павлов [Artyom Pavlov] <[email protected]>

Author: nstilt1

chacha20/src/backends/avx2.rs

@@ -1,9 +1,9 @@
 #![allow(unsafe_op_in_unsafe_fn)]
-use crate::Rounds;
+use crate::{Rounds, Variant};
 use core::marker::PhantomData;
 
 #[cfg(feature = "rng")]
-use crate::{ChaChaCore, Variant};
+use crate::ChaChaCore;
 
 #[cfg(feature = "cipher")]
 use crate::{chacha::Block, STATE_WORDS};
@@ -27,10 +27,11 @@ const N: usize = PAR_BLOCKS / 2;
 #[inline]
 #[target_feature(enable = "avx2")]
 #[cfg(feature = "cipher")]
-pub(crate) unsafe fn inner<R, F>(state: &mut [u32; STATE_WORDS], f: F)
+pub(crate) unsafe fn inner<R, F, V>(state: &mut [u32; STATE_WORDS], f: F)
 where
     R: Rounds,
     F: StreamCipherClosure<BlockSize = U64>,
+    V: Variant,
 {
     let state_ptr = state.as_ptr() as *const __m128i;
     let v = [
@@ -39,13 +40,21 @@ where
         _mm256_broadcastsi128_si256(_mm_loadu_si128(state_ptr.add(2))),
     ];
     let mut c = _mm256_broadcastsi128_si256(_mm_loadu_si128(state_ptr.add(3)));
-    c = _mm256_add_epi32(c, _mm256_set_epi32(0, 0, 0, 1, 0, 0, 0, 0));
+    c = match size_of::<V::Counter>() {
+        4 => _mm256_add_epi32(c, _mm256_set_epi32(0, 0, 0, 1, 0, 0, 0, 0)),
+        8 => _mm256_add_epi64(c, _mm256_set_epi64x(0, 1, 0, 0)),
+        _ => unreachable!()
+    };
     let mut ctr = [c; N];
     for i in 0..N {
         ctr[i] = c;
-        c = _mm256_add_epi32(c, _mm256_set_epi32(0, 0, 0, 2, 0, 0, 0, 2));
+        c = match size_of::<V::Counter>() {
+            4 => _mm256_add_epi32(c, _mm256_set_epi32(0, 0, 0, 2, 0, 0, 0, 2)),
+            8 => _mm256_add_epi64(c, _mm256_set_epi64x(0, 2, 0, 2)),
+            _ => unreachable!(),
+        };
     }
-    let mut backend = Backend::<R> {
+    let mut backend = Backend::<R, V> {
         v,
         ctr,
         _pd: PhantomData,
@@ -54,6 +63,11 @@ where
     f.call(&mut backend);
 
     state[12] = _mm256_extract_epi32(backend.ctr[0], 0) as u32;
+    match size_of::<V::Counter>() {
+        4 => {},
+        8 => state[13] = _mm256_extract_epi32(backend.ctr[0], 1) as u32,
+        _ => unreachable!()
+    }
 }
 
 #[inline]
@@ -71,13 +85,13 @@ where
         _mm256_broadcastsi128_si256(_mm_loadu_si128(state_ptr.add(2))),
     ];
     let mut c = _mm256_broadcastsi128_si256(_mm_loadu_si128(state_ptr.add(3)));
-    c = _mm256_add_epi32(c, _mm256_set_epi32(0, 0, 0, 1, 0, 0, 0, 0));
+    c = _mm256_add_epi64(c, _mm256_set_epi64x(0, 1, 0, 0));
     let mut ctr = [c; N];
     for i in 0..N {
         ctr[i] = c;
-        c = _mm256_add_epi32(c, _mm256_set_epi32(0, 0, 0, 2, 0, 0, 0, 2));
+        c = _mm256_add_epi64(c, _mm256_set_epi64x(0, 2, 0, 2));
     }
-    let mut backend = Backend::<R> {
+    let mut backend = Backend::<R, V> {
         v,
         ctr,
         _pd: PhantomData,
@@ -86,32 +100,37 @@ where
     backend.rng_gen_par_ks_blocks(buffer);
 
     core.state[12] = _mm256_extract_epi32(backend.ctr[0], 0) as u32;
+    core.state[13] = _mm256_extract_epi32(backend.ctr[0], 1) as u32;
 }
 
-struct Backend<R: Rounds> {
+struct Backend<R: Rounds, V: Variant> {
     v: [__m256i; 3],
     ctr: [__m256i; N],
-    _pd: PhantomData<R>,
+    _pd: PhantomData<(R, V)>,
 }
 
 #[cfg(feature = "cipher")]
-impl<R: Rounds> BlockSizeUser for Backend<R> {
+impl<R: Rounds, V: Variant> BlockSizeUser for Backend<R, V> {
     type BlockSize = U64;
 }
 
 #[cfg(feature = "cipher")]
-impl<R: Rounds> ParBlocksSizeUser for Backend<R> {
+impl<R: Rounds, V: Variant> ParBlocksSizeUser for Backend<R, V> {
     type ParBlocksSize = U4;
 }
 
 #[cfg(feature = "cipher")]
-impl<R: Rounds> StreamCipherBackend for Backend<R> {
+impl<R: Rounds, V: Variant> StreamCipherBackend for Backend<R, V> {
     #[inline(always)]
     fn gen_ks_block(&mut self, block: &mut Block) {
         unsafe {
             let res = rounds::<R>(&self.v, &self.ctr);
             for c in self.ctr.iter_mut() {
-                *c = _mm256_add_epi32(*c, _mm256_set_epi32(0, 0, 0, 1, 0, 0, 0, 1));
+                *c = match size_of::<V::Counter>() {
+                    4 => _mm256_add_epi32(*c, _mm256_set_epi32(0, 0, 0, 1, 0, 0, 0, 1)),
+                    8 => _mm256_add_epi64(*c, _mm256_set_epi64x(0, 1, 0, 1)),
+                    _ => unreachable!()
+                };
             }
 
             let res0: [__m128i; 8] = core::mem::transmute(res[0]);
@@ -130,7 +149,11 @@ impl<R: Rounds> StreamCipherBackend for Backend<R> {
 
             let pb = PAR_BLOCKS as i32;
             for c in self.ctr.iter_mut() {
-                *c = _mm256_add_epi32(*c, _mm256_set_epi32(0, 0, 0, pb, 0, 0, 0, pb));
+                *c = match size_of::<V::Counter>() {
+                    4 => _mm256_add_epi32(*c, _mm256_set_epi32(0, 0, 0, pb, 0, 0, 0, pb)),
+                    8 => _mm256_add_epi64(*c, _mm256_set_epi64x(0, pb as i64, 0, pb as i64)),
+                    _ => unreachable!()
+                }
             }
 
             let mut block_ptr = blocks.as_mut_ptr() as *mut __m128i;
@@ -147,15 +170,15 @@ impl<R: Rounds> StreamCipherBackend for Backend<R> {
 }
 
 #[cfg(feature = "rng")]
-impl<R: Rounds> Backend<R> {
+impl<R: Rounds, V: Variant> Backend<R, V> {
     #[inline(always)]
     fn rng_gen_par_ks_blocks(&mut self, blocks: &mut [u32; 64]) {
         unsafe {
             let vs = rounds::<R>(&self.v, &self.ctr);
 
             let pb = PAR_BLOCKS as i32;
             for c in self.ctr.iter_mut() {
-                *c = _mm256_add_epi32(*c, _mm256_set_epi32(0, 0, 0, pb, 0, 0, 0, pb));
+                *c = _mm256_add_epi64(*c, _mm256_set_epi64x(0, pb as i64, 0, pb as i64));
             }
 
             let mut block_ptr = blocks.as_mut_ptr() as *mut __m128i;

chacha20/src/backends/neon.rs

@@ -4,11 +4,11 @@
 //! Adapted from the Crypto++ `chacha_simd` implementation by Jack Lloyd and
 //! Jeffrey Walton (public domain).
 
-use crate::{Rounds, STATE_WORDS};
+use crate::{Rounds, STATE_WORDS, Variant};
 use core::{arch::aarch64::*, marker::PhantomData};
 
 #[cfg(feature = "rand_core")]
-use crate::{ChaChaCore, Variant};
+use crate::ChaChaCore;
 
 #[cfg(feature = "cipher")]
 use crate::chacha::Block;
@@ -19,13 +19,26 @@ use cipher::{
     consts::{U4, U64},
 };
 
-struct Backend<R: Rounds> {
+struct Backend<R: Rounds, V: Variant> {
     state: [uint32x4_t; 4],
     ctrs: [uint32x4_t; 4],
-    _pd: PhantomData<R>,
+    _pd: PhantomData<(R, V)>,
 }
 
-impl<R: Rounds> Backend<R> {
+macro_rules! add_counter {
+    ($a:expr, $b:expr, $variant:ty) => {
+        match size_of::<<$variant>::Counter>() {
+            4 => vaddq_u32($a, $b),
+            8 => vreinterpretq_u32_u64(vaddq_u64(
+                vreinterpretq_u64_u32($a),
+                vreinterpretq_u64_u32($b),
+            )),
+            _ => unreachable!(),
+        }
+    };
+}
+
+impl<R: Rounds, V: Variant> Backend<R, V> {
     #[inline]
     unsafe fn new(state: &mut [u32; STATE_WORDS]) -> Self {
         let state = [
@@ -40,7 +53,7 @@ impl<R: Rounds> Backend<R> {
             vld1q_u32([3, 0, 0, 0].as_ptr()),
             vld1q_u32([4, 0, 0, 0].as_ptr()),
         ];
-        Backend::<R> {
+        Backend::<R, V> {
             state,
             ctrs,
             _pd: PhantomData,
@@ -51,16 +64,24 @@ impl<R: Rounds> Backend<R> {
 #[inline]
 #[cfg(feature = "cipher")]
 #[target_feature(enable = "neon")]
-pub(crate) unsafe fn inner<R, F>(state: &mut [u32; STATE_WORDS], f: F)
+pub(crate) unsafe fn inner<R, F, V>(state: &mut [u32; STATE_WORDS], f: F)
 where
     R: Rounds,
     F: StreamCipherClosure<BlockSize = U64>,
+    V: Variant,
 {
-    let mut backend = Backend::<R>::new(state);
+    let mut backend = Backend::<R, V>::new(state);
 
     f.call(&mut backend);
 
-    vst1q_u32(state.as_mut_ptr().offset(12), backend.state[3]);
+    match size_of::<V::Counter>() {
+        4 => state[12] = vgetq_lane_u32(backend.state[3], 0),
+        8 => vst1q_u64(
+            state.as_mut_ptr().offset(12) as *mut u64,
+            vreinterpretq_u64_u32(backend.state[3]),
+        ),
+        _ => unreachable!(),
+    }
 }
 
 #[inline]
@@ -73,19 +94,22 @@ where
     R: Rounds,
     V: Variant,
 {
-    let mut backend = Backend::<R>::new(&mut core.state);
+    let mut backend = Backend::<R, V>::new(&mut core.state);
 
     backend.write_par_ks_blocks(buffer);
 
-    vst1q_u32(core.state.as_mut_ptr().offset(12), backend.state[3]);
+    vst1q_u64(
+        core.state.as_mut_ptr().offset(12) as *mut u64,
+        vreinterpretq_u64_u32(backend.state[3]),
+    );
 }
 
 #[cfg(feature = "cipher")]
-impl<R: Rounds> BlockSizeUser for Backend<R> {
+impl<R: Rounds, V: Variant> BlockSizeUser for Backend<R, V> {
     type BlockSize = U64;
 }
 #[cfg(feature = "cipher")]
-impl<R: Rounds> ParBlocksSizeUser for Backend<R> {
+impl<R: Rounds, V: Variant> ParBlocksSizeUser for Backend<R, V> {
     type ParBlocksSize = U4;
 }
 
@@ -97,15 +121,15 @@ macro_rules! add_assign_vec {
 }
 
 #[cfg(feature = "cipher")]
-impl<R: Rounds> StreamCipherBackend for Backend<R> {
+impl<R: Rounds, V: Variant> StreamCipherBackend for Backend<R, V> {
     #[inline(always)]
     fn gen_ks_block(&mut self, block: &mut Block) {
         let state3 = self.state[3];
         let mut par = ParBlocks::<Self>::default();
         self.gen_par_ks_blocks(&mut par);
         *block = par[0];
         unsafe {
-            self.state[3] = vaddq_u32(state3, vld1q_u32([1, 0, 0, 0].as_ptr()));
+            self.state[3] = add_counter!(state3, vld1q_u32([1, 0, 0, 0].as_ptr()), V);
         }
     }
 
@@ -118,19 +142,19 @@ impl<R: Rounds> StreamCipherBackend for Backend<R> {
                     self.state[0],
                     self.state[1],
                     self.state[2],
-                    vaddq_u32(self.state[3], self.ctrs[0]),
+                    add_counter!(self.state[3], self.ctrs[0], V),
                 ],
                 [
                     self.state[0],
                     self.state[1],
                     self.state[2],
-                    vaddq_u32(self.state[3], self.ctrs[1]),
+                    add_counter!(self.state[3], self.ctrs[1], V),
                 ],
                 [
                     self.state[0],
                     self.state[1],
                     self.state[2],
-                    vaddq_u32(self.state[3], self.ctrs[2]),
+                    add_counter!(self.state[3], self.ctrs[2], V),
                 ],
             ];
 
@@ -140,11 +164,16 @@ impl<R: Rounds> StreamCipherBackend for Backend<R> {
 
             for block in 0..4 {
                 // add state to block
-                for state_row in 0..4 {
+                for state_row in 0..3 {
                     add_assign_vec!(blocks[block][state_row], self.state[state_row]);
                 }
                 if block > 0 {
-                    blocks[block][3] = vaddq_u32(blocks[block][3], self.ctrs[block - 1]);
+                    add_assign_vec!(
+                        blocks[block][3],
+                        add_counter!(self.state[3], self.ctrs[block - 1], V)
+                    );
+                } else {
+                    add_assign_vec!(blocks[block][3], self.state[3]);
                 }
                 // write blocks to dest
                 for state_row in 0..4 {
@@ -154,7 +183,7 @@ impl<R: Rounds> StreamCipherBackend for Backend<R> {
                     );
                 }
             }
-            self.state[3] = vaddq_u32(self.state[3], self.ctrs[3]);
+            self.state[3] = add_counter!(self.state[3], self.ctrs[3], V);
         }
     }
 }
@@ -180,7 +209,7 @@ macro_rules! extract {
     };
 }
 
-impl<R: Rounds> Backend<R> {
+impl<R: Rounds, V: Variant> Backend<R, V> {
     #[inline(always)]
     /// Generates `num_blocks` blocks and blindly writes them to `dest_ptr`
     ///
@@ -197,19 +226,19 @@ impl<R: Rounds> Backend<R> {
                 self.state[0],
                 self.state[1],
                 self.state[2],
-                vaddq_u32(self.state[3], self.ctrs[0]),
+                add_counter!(self.state[3], self.ctrs[0], V),
             ],
             [
                 self.state[0],
                 self.state[1],
                 self.state[2],
-                vaddq_u32(self.state[3], self.ctrs[1]),
+                add_counter!(self.state[3], self.ctrs[1], V),
             ],
             [
                 self.state[0],
                 self.state[1],
                 self.state[2],
-                vaddq_u32(self.state[3], self.ctrs[2]),
+                add_counter!(self.state[3], self.ctrs[2], V),
             ],
         ];
 
@@ -220,11 +249,16 @@ impl<R: Rounds> Backend<R> {
         let mut dest_ptr = buffer.as_mut_ptr() as *mut u8;
         for block in 0..4 {
             // add state to block
-            for state_row in 0..4 {
+            for state_row in 0..3 {
                 add_assign_vec!(blocks[block][state_row], self.state[state_row]);
             }
             if block > 0 {
-                blocks[block][3] = vaddq_u32(blocks[block][3], self.ctrs[block - 1]);
+                add_assign_vec!(
+                    blocks[block][3],
+                    add_counter!(self.state[3], self.ctrs[block - 1], V)
+                );
+            } else {
+                add_assign_vec!(blocks[block][3], self.state[3]);
             }
             // write blocks to buffer
             for state_row in 0..4 {
@@ -235,7 +269,7 @@ impl<R: Rounds> Backend<R> {
             }
             dest_ptr = dest_ptr.add(64);
         }
-        self.state[3] = vaddq_u32(self.state[3], self.ctrs[3]);
+        self.state[3] = add_counter!(self.state[3], self.ctrs[3], V);
     }
 }