salsa20: add type param for key length (#432)

Author: micolous

salsa20/src/backends/soft.rs

@@ -7,17 +7,17 @@ use cipher::{
     consts::{U1, U64},
 };
 
-pub(crate) struct Backend<'a, R: Unsigned>(pub(crate) &'a mut SalsaCore<R>);
+pub(crate) struct Backend<'a, R: Unsigned, KeySize>(pub(crate) &'a mut SalsaCore<R, KeySize>);
 
-impl<R: Unsigned> BlockSizeUser for Backend<'_, R> {
+impl<R: Unsigned, KeySize> BlockSizeUser for Backend<'_, R, KeySize> {
     type BlockSize = U64;
 }
 
-impl<R: Unsigned> ParBlocksSizeUser for Backend<'_, R> {
+impl<R: Unsigned, KeySize> ParBlocksSizeUser for Backend<'_, R, KeySize> {
     type ParBlocksSize = U1;
 }
 
-impl<R: Unsigned> StreamCipherBackend for Backend<'_, R> {
+impl<R: Unsigned, KeySize> StreamCipherBackend for Backend<'_, R, KeySize> {
     #[inline(always)]
     fn gen_ks_block(&mut self, block: &mut Block<Self>) {
         let res = run_rounds::<R>(&self.0.state);

salsa20/src/backends/sse2.rs

@@ -16,7 +16,7 @@ use core::arch::x86_64::*;
 
 #[inline]
 #[target_feature(enable = "sse2")]
-pub(crate) unsafe fn inner<R, F>(state: &mut [u32; STATE_WORDS], f: F)
+pub(crate) unsafe fn inner<R, F, KeySize>(state: &mut [u32; STATE_WORDS], f: F)
 where
     R: Unsigned,
     F: StreamCipherClosure<BlockSize = U64>,
@@ -37,9 +37,10 @@ where
         f.call(&mut backend);
         state[8] = _mm_cvtsi128_si32(backend.v[2]) as u32;
     } else {
-        f.call(&mut SoftBackend(&mut SalsaCore::<R> {
+        f.call(&mut SoftBackend(&mut SalsaCore::<R, KeySize> {
             state: *state,
             rounds: PhantomData,
+            key_size: PhantomData,
         }));
     }
 }

salsa20/src/lib.rs

@@ -80,7 +80,7 @@ pub use cipher;
 use cipher::{
     Block, BlockSizeUser, IvSizeUser, KeyIvInit, KeySizeUser, StreamCipherClosure,
     StreamCipherCore, StreamCipherCoreWrapper, StreamCipherSeekCore,
-    array::{Array, typenum::Unsigned},
+    array::{Array, ArraySize, typenum::Unsigned},
     consts::{U4, U6, U8, U10, U24, U32, U64},
 };
 use core::marker::PhantomData;
@@ -95,18 +95,18 @@ pub use xsalsa::{XSalsa8, XSalsa12, XSalsa20, XSalsaCore, hsalsa};
 
 /// Salsa20/8 stream cipher
 /// (reduced-round variant of Salsa20 with 8 rounds, *not recommended*)
-pub type Salsa8 = StreamCipherCoreWrapper<SalsaCore<U4>>;
+pub type Salsa8 = StreamCipherCoreWrapper<SalsaCore<U4, U32>>;
 
 /// Salsa20/12 stream cipher
 /// (reduced-round variant of Salsa20 with 12 rounds, *not recommended*)
-pub type Salsa12 = StreamCipherCoreWrapper<SalsaCore<U6>>;
+pub type Salsa12 = StreamCipherCoreWrapper<SalsaCore<U6, U32>>;
 
 /// Salsa20/20 stream cipher
 /// (20 rounds; **recommended**)
-pub type Salsa20 = StreamCipherCoreWrapper<SalsaCore<U10>>;
+pub type Salsa20 = StreamCipherCoreWrapper<SalsaCore<U10, U32>>;
 
 /// Key type used by all Salsa variants and [`XSalsa20`].
-pub type Key = Array<u8, U32>;
+pub type Key<KeySize> = Array<u8, KeySize>;
 
 /// Nonce type used by all Salsa variants.
 pub type Nonce = Array<u8, U8>;
@@ -121,14 +121,16 @@ const STATE_WORDS: usize = 16;
 const CONSTANTS: [u32; 4] = [0x6170_7865, 0x3320_646e, 0x7962_2d32, 0x6b20_6574];
 
 /// The Salsa20 core function.
-pub struct SalsaCore<R: Unsigned> {
+pub struct SalsaCore<R: Unsigned, KeySize = U32> {
     /// Internal state of the core function
     state: [u32; STATE_WORDS],
     /// Number of rounds to perform
     rounds: PhantomData<R>,
+    /// Key size
+    key_size: PhantomData<KeySize>,
 }
 
-impl<R: Unsigned> SalsaCore<R> {
+impl<R: Unsigned, KeySize> SalsaCore<R, KeySize> {
     /// Create new Salsa core from raw state.
     ///
     /// This method is mainly intended for the `scrypt` crate.
@@ -137,24 +139,29 @@ impl<R: Unsigned> SalsaCore<R> {
         Self {
             state,
             rounds: PhantomData,
+            key_size: PhantomData,
         }
     }
 }
 
-impl<R: Unsigned> KeySizeUser for SalsaCore<R> {
-    type KeySize = U32;
+impl<R: Unsigned, KeySize> KeySizeUser for SalsaCore<R, KeySize>
+where
+    KeySize: ArraySize,
+{
+    type KeySize = KeySize;
 }
 
-impl<R: Unsigned> IvSizeUser for SalsaCore<R> {
+impl<R: Unsigned, KeySize> IvSizeUser for SalsaCore<R, KeySize> {
     type IvSize = U8;
 }
 
-impl<R: Unsigned> BlockSizeUser for SalsaCore<R> {
+impl<R: Unsigned, KeySize> BlockSizeUser for SalsaCore<R, KeySize> {
     type BlockSize = U64;
 }
 
-impl<R: Unsigned> KeyIvInit for SalsaCore<R> {
-    fn new(key: &Key, iv: &Nonce) -> Self {
+impl<R: Unsigned> KeyIvInit for SalsaCore<R, U32>
+{
+    fn new(key: &Key<U32>, iv: &Nonce) -> Self {
         let mut state = [0u32; STATE_WORDS];
         state[0] = CONSTANTS[0];
 
@@ -192,11 +199,12 @@ impl<R: Unsigned> KeyIvInit for SalsaCore<R> {
         Self {
             state,
             rounds: PhantomData,
+            key_size: PhantomData,
         }
     }
 }
 
-impl<R: Unsigned> StreamCipherCore for SalsaCore<R> {
+impl<R: Unsigned, KeySize> StreamCipherCore for SalsaCore<R, KeySize> {
     #[inline(always)]
     fn remaining_blocks(&self) -> Option<usize> {
         let rem = u64::MAX - self.get_block_pos();
@@ -206,7 +214,7 @@ impl<R: Unsigned> StreamCipherCore for SalsaCore<R> {
         cfg_if! {
             if #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] {
                 unsafe {
-                    backends::sse2::inner::<R, _>(&mut self.state, f);
+                    backends::sse2::inner::<R, _, KeySize>(&mut self.state, f);
                 }
             } else {
                 f.call(&mut backends::soft::Backend(self));
@@ -215,7 +223,7 @@ impl<R: Unsigned> StreamCipherCore for SalsaCore<R> {
     }
 }
 
-impl<R: Unsigned> StreamCipherSeekCore for SalsaCore<R> {
+impl<R: Unsigned, KeySize> StreamCipherSeekCore for SalsaCore<R, KeySize> {
     type Counter = u64;
 
     #[inline(always)]

salsa20/src/xsalsa.rs

@@ -25,7 +25,7 @@ pub type XSalsa12 = StreamCipherCoreWrapper<XSalsaCore<U6>>;
 pub type XSalsa8 = StreamCipherCoreWrapper<XSalsaCore<U4>>;
 
 /// The XSalsa core function.
-pub struct XSalsaCore<R: Unsigned>(SalsaCore<R>);
+pub struct XSalsaCore<R: Unsigned>(SalsaCore<R, U32>);
 
 impl<R: Unsigned> KeySizeUser for XSalsaCore<R> {
     type KeySize = U32;
@@ -41,7 +41,7 @@ impl<R: Unsigned> BlockSizeUser for XSalsaCore<R> {
 
 impl<R: Unsigned> KeyIvInit for XSalsaCore<R> {
     #[inline]
-    fn new(key: &Key, iv: &XNonce) -> Self {
+    fn new(key: &Key<U32>, iv: &XNonce) -> Self {
         let subkey = hsalsa::<R>(key, iv[..16].try_into().unwrap());
         let mut padded_iv = Nonce::default();
         padded_iv.copy_from_slice(&iv[16..]);
@@ -89,7 +89,7 @@ impl<R: Unsigned> ZeroizeOnDrop for XSalsaCore<R> {}
 /// - Nonce (`u32` x 4)
 ///
 /// It produces 256-bits of output suitable for use as a Salsa20 key
-pub fn hsalsa<R: Unsigned>(key: &Key, input: &Array<u8, U16>) -> Array<u8, U32> {
+pub fn hsalsa<R: Unsigned>(key: &Key<U32>, input: &Array<u8, U16>) -> Array<u8, U32> {
     #[inline(always)]
     fn to_u32(chunk: &[u8]) -> u32 {
         u32::from_le_bytes(chunk.try_into().unwrap())