vru_noise/
hash.rs

1use core::ops::Mul;
2
3use {
4    digest::{Update, core_api::BlockSizeUser, FixedOutput},
5    aead::KeySizeUser,
6    generic_array::{
7        GenericArray, ArrayLength,
8        sequence::GenericSequence,
9        typenum::{self, Unsigned},
10    },
11    digest::OutputSizeUser,
12    hkdf::HmacImpl,
13};
14
15pub trait MixHash
16where
17    Self: Sized,
18{
19    type L: ArrayLength<u8>;
20    type B: ArrayLength<u8>;
21
22    fn init(data: &[u8]) -> GenericArray<u8, Self::L>;
23
24    fn mix_hash(hash: GenericArray<u8, Self::L>, data: &[u8]) -> GenericArray<u8, Self::L>;
25
26    fn mix_parts(hash: GenericArray<u8, Self::L>, parts: &[&[u8]]) -> GenericArray<u8, Self::L>;
27}
28
29impl<D> MixHash for D
30where
31    D: BlockSizeUser + Update + FixedOutput + Default,
32{
33    type L = D::OutputSize;
34    type B = D::BlockSize;
35
36    fn init(data: &[u8]) -> GenericArray<u8, Self::L> {
37        D::default().chain(data).finalize_fixed()
38    }
39
40    fn mix_hash(hash: GenericArray<u8, Self::L>, data: &[u8]) -> GenericArray<u8, Self::L> {
41        D::default().chain(hash).chain(data).finalize_fixed()
42    }
43
44    fn mix_parts(hash: GenericArray<u8, Self::L>, parts: &[&[u8]]) -> GenericArray<u8, Self::L> {
45        let mut d = D::default().chain(&hash);
46        for &part in parts {
47            d.update(part);
48        }
49        d.finalize_fixed()
50    }
51}
52
53pub trait HkdfSplit<N>
54where
55    N: ArrayLength<GenericArray<u8, Self::L>>,
56{
57    type L: ArrayLength<u8>;
58
59    fn hkdf_split(salt: Option<&[u8]>, ikm: &[u8]) -> GenericArray<GenericArray<u8, Self::L>, N>;
60}
61
62impl<D, I, N> HkdfSplit<N> for (D, I)
63where
64    I: HmacImpl<D>,
65    D: OutputSizeUser,
66    N: ArrayLength<GenericArray<u8, D::OutputSize>>,
67    D::OutputSize: Mul<N>,
68    <D::OutputSize as Mul<N>>::Output: ArrayLength<u8>,
69{
70    type L = D::OutputSize;
71
72    fn hkdf_split(salt: Option<&[u8]>, ikm: &[u8]) -> GenericArray<GenericArray<u8, Self::L>, N> {
73        use hkdf::Hkdf;
74
75        let hkdf = Hkdf::<D, I>::new(salt, ikm);
76        let mut okm: GenericArray<u8, <Self::L as Mul<N>>::Output> = GenericArray::default();
77        hkdf.expand(&[], okm.as_mut()).unwrap();
78        let l = <Self::L as Unsigned>::USIZE;
79        GenericArray::generate(|i| {
80            let mut s = GenericArray::default();
81            s.as_mut_slice()
82                .clone_from_slice(&okm[(l * i)..(l * (i + 1))]);
83            s
84        })
85    }
86}
87
88pub trait HkdfSplitExt<A>
89where
90    A: KeySizeUser,
91{
92    type L: ArrayLength<u8>;
93
94    fn split_final(
95        chaining_key: &[u8],
96        data: &[u8],
97    ) -> (GenericArray<u8, A::KeySize>, GenericArray<u8, A::KeySize>);
98
99    fn split_2(
100        chaining_key: &[u8],
101        data: &[u8],
102    ) -> (GenericArray<u8, Self::L>, GenericArray<u8, A::KeySize>);
103
104    fn split_3(
105        chaining_key: &[u8],
106        data: &[u8],
107    ) -> (
108        GenericArray<u8, Self::L>,
109        GenericArray<u8, Self::L>,
110        GenericArray<u8, A::KeySize>,
111    );
112}
113
114fn truncate<A>(chaining_key: &[u8]) -> GenericArray<u8, A::KeySize>
115where
116    A: KeySizeUser,
117{
118    let input_length = chaining_key.len();
119    let output_length = A::KeySize::USIZE;
120    assert!(output_length <= input_length);
121
122    let mut a = GenericArray::default();
123    a[..output_length].clone_from_slice(&chaining_key[..output_length]);
124    a
125}
126
127impl<A, T> HkdfSplitExt<A> for T
128where
129    A: KeySizeUser,
130    T: HkdfSplit<typenum::U2> + HkdfSplit<typenum::U3, L = <Self as HkdfSplit<typenum::U2>>::L>,
131{
132    type L = <T as HkdfSplit<typenum::U2>>::L;
133
134    fn split_final(
135        chaining_key: &[u8],
136        data: &[u8],
137    ) -> (GenericArray<u8, A::KeySize>, GenericArray<u8, A::KeySize>) {
138        let keys = Self::hkdf_split(Some(chaining_key), data);
139        let [send_key, receive_key]: [_; 2] = keys.into();
140        (
141            truncate::<A>(send_key.as_ref()),
142            truncate::<A>(receive_key.as_ref()),
143        )
144    }
145
146    fn split_2(
147        chaining_key: &[u8],
148        data: &[u8],
149    ) -> (GenericArray<u8, Self::L>, GenericArray<u8, A::KeySize>) {
150        let keys = Self::hkdf_split(Some(chaining_key), data);
151        let [chaining_key, key]: [_; 2] = keys.into();
152        (chaining_key, truncate::<A>(key.as_ref()))
153    }
154
155    fn split_3(
156        chaining_key: &[u8],
157        data: &[u8],
158    ) -> (
159        GenericArray<u8, Self::L>,
160        GenericArray<u8, Self::L>,
161        GenericArray<u8, A::KeySize>,
162    ) {
163        let keys = Self::hkdf_split(Some(chaining_key), data);
164        let [chaining_key, middle, key]: [_; 3] = keys.into();
165        (chaining_key, middle, truncate::<A>(key.as_ref()))
166    }
167}