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}