A whimsical STROBE based encryption protocol
1use core::marker::PhantomData;
2
3use aead::{
4 Buffer,
5 consts::{U32, U64},
6};
7use alloc::boxed::Box;
8use hybrid_array::{Array, SliceExt, typenum::Unsigned};
9use ml_kem::{Encapsulate, Kem, KeyExport, ParameterSet, SharedKey, TryKeyInit, kem::Decapsulate};
10use rand_core::CryptoRng;
11use wharrgarbl_neko::{NekoSec, NekoState};
12use zeroize::Zeroize;
13
14use crate::{
15 Role, WHARRGHARBL_PROTO,
16 transport::{AeadNeko, AeadTransport},
17};
18
19pub struct ClientHandshake<S: NekoSec, K: Kem + ParameterSet> {
20 kem_sec: PhantomData<K>,
21 neko: NekoState<S>,
22 decap: Option<Box<K::DecapsulationKey>>,
23}
24
25impl<S, K> ClientHandshake<S, K>
26where
27 S: NekoSec,
28 K: Kem<SharedKeySize = U32> + ParameterSet,
29 K::DecapsulationKey: Decapsulate,
30{
31 pub fn new(psk: Option<&[u8; 32]>) -> Self {
32 let mut neko = NekoState::<S>::new(WHARRGHARBL_PROTO.as_bytes());
33
34 if let Some(psk) = psk.and_then(|psk| psk.as_hybrid_array()) {
35 neko.key(psk);
36 }
37
38 neko.ad(&K::K::to_u8().to_le_bytes());
39 neko.ad(&S::to_bytes());
40 neko.ratchet();
41
42 Self {
43 kem_sec: PhantomData,
44 neko,
45 decap: None,
46 }
47 }
48
49 pub fn send(&mut self, rng: &mut impl CryptoRng, buf: &mut dyn Buffer) -> aead::Result<()> {
50 let (decap, encap) = K::generate_keypair_from_rng(rng);
51
52 self.decap = Some(Box::new(decap));
53
54 buf.extend_from_slice(&encap.to_bytes())?;
55
56 self.neko.cleartext(buf.as_ref());
57
58 buf.extend_from_slice(&self.neko.create_mac())?;
59
60 Ok(())
61 }
62
63 pub fn receive(&mut self, ciphertext: &[u8]) -> aead::Result<AeadTransport<S>> {
64 let decap = self.decap.as_ref().ok_or(aead::Error)?;
65
66 let tag = ciphertext
67 .len()
68 .checked_sub(<AeadNeko<S> as aead::AeadCore>::TagSize::to_usize())
69 .ok_or(aead::Error)?;
70
71 let (ciphertext, tag) = ciphertext.split_at(tag);
72
73 let tag: aead::Tag<AeadNeko<S>> = tag.try_into().unwrap();
74
75 self.neko.cleartext(ciphertext);
76 self.neko.verify_mac(&tag)?;
77
78 let shared = decap
79 .decapsulate_slice(ciphertext)
80 .map_err(|_| aead::Error)?;
81
82 Ok(self.finish(shared))
83 }
84
85 fn finish(&mut self, mut shared: SharedKey) -> AeadTransport<S> {
86 self.neko.key(&shared);
87 shared.zeroize();
88
89 self.neko.ratchet();
90
91 let mut entropy: Array<u8, U64> = Default::default();
92
93 self.neko.prf(&mut entropy);
94
95 let (key, rest) = entropy.split();
96 let (outbound, inbound) = rest.split();
97
98 entropy.zeroize();
99
100 AeadTransport::new(key, outbound, inbound, Role::Sender)
101 }
102}
103
104pub struct ServerHandshake<S: NekoSec, K: Kem + ParameterSet> {
105 kem_sec: PhantomData<K>,
106 neko: NekoState<S>,
107}
108
109impl<S, K> ServerHandshake<S, K>
110where
111 S: NekoSec,
112 K: Kem<SharedKeySize = U32> + ParameterSet,
113{
114 pub fn new(psk: Option<&[u8; 32]>) -> Self {
115 let mut neko = NekoState::<S>::new(WHARRGHARBL_PROTO.as_bytes());
116
117 if let Some(psk) = psk.and_then(|psk| psk.as_hybrid_array()) {
118 neko.key(psk);
119 }
120
121 neko.ad(&K::K::to_u8().to_le_bytes());
122 neko.ad(&S::to_bytes());
123 neko.ratchet();
124
125 Self {
126 kem_sec: PhantomData,
127 neko,
128 }
129 }
130
131 pub fn respond(
132 &mut self,
133 rng: &mut impl CryptoRng,
134 buf: &mut dyn Buffer,
135 ) -> aead::Result<AeadTransport<S>> {
136 let slice = buf.as_ref();
137
138 let tag = slice
139 .len()
140 .checked_sub(<AeadNeko<S> as aead::AeadCore>::TagSize::to_usize())
141 .ok_or(aead::Error)?;
142
143 let (encap, tag) = buf.as_ref().split_at(tag);
144
145 let tag: aead::Tag<AeadNeko<S>> = tag.try_into().unwrap();
146
147 self.neko.cleartext(encap);
148 self.neko.verify_mac(&tag)?;
149
150 let encap = K::EncapsulationKey::new_from_slice(encap).map_err(|_| aead::Error)?;
151
152 let (cipher, shared) = encap.encapsulate_with_rng(rng);
153
154 buf.truncate(0);
155
156 buf.extend_from_slice(cipher.as_ref())?;
157 self.neko.cleartext(cipher.as_ref());
158 buf.extend_from_slice(&self.neko.create_mac())?;
159
160 Ok(self.finish(shared))
161 }
162
163 fn finish(&mut self, mut shared: SharedKey) -> AeadTransport<S> {
164 self.neko.key(&shared);
165 shared.zeroize();
166
167 self.neko.ratchet();
168
169 let mut entropy: Array<u8, U64> = Default::default();
170
171 self.neko.prf(&mut entropy);
172
173 let (key, rest) = entropy.split();
174 let (outbound, inbound) = rest.split();
175
176 entropy.zeroize();
177
178 AeadTransport::new(key, outbound, inbound, Role::Receiver)
179 }
180}
181
182#[cfg(test)]
183mod tests {
184 use ml_kem::{MlKem512, MlKem768};
185 use wharrgarbl_neko::{Neko128, Neko256};
186
187 use crate::utils::BufferSlice;
188
189 use super::*;
190
191 #[test]
192 fn handshake_protocol_happy_path() -> aead::Result<()> {
193 let psk: [u8; 32] = [
194 31, 48, 29, 177, 88, 236, 186, 84, 65, 51, 214, 243, 174, 24, 45, 101, 229, 129, 62,
195 132, 45, 174, 183, 65, 89, 73, 107, 177, 77, 90, 164, 251,
196 ];
197
198 let mut alice = ClientHandshake::<Neko128, MlKem512>::new(Some(&psk));
199 let mut bob = ServerHandshake::<Neko128, MlKem512>::new(Some(&psk));
200
201 // BufferSlice acts as our transport across the webz
202 let mut buf = alloc::vec![0u8; 2048];
203 let mut buf = BufferSlice::new(&mut buf);
204
205 let mut rng = rand_core::UnwrapErr(getrandom::SysRng);
206
207 alice.send(&mut rng, &mut buf).unwrap();
208
209 // Pretend to send ek across the webz: client -> server
210 let bob = bob.respond(&mut rng, &mut buf).unwrap();
211
212 // Pretend to send ciphertext across the webz: server -> client
213 let alice = alice.receive(buf.as_ref()).unwrap();
214
215 assert_eq!(alice.aead.key, bob.aead.key);
216
217 // Both AeadStates have derived base nonces for each context.
218 // Inbound context nonces will not match Outbound context nonces.
219 assert_eq!(alice.trump, bob.trump);
220 assert_eq!(alice.epstein, bob.epstein);
221 assert_ne!(alice.trump, alice.epstein);
222 assert_ne!(bob.trump, bob.epstein);
223
224 Ok(())
225 }
226
227 #[test]
228 #[should_panic]
229 fn handshake_fails_with_neko_security_mismatch() {
230 let psk: [u8; 32] = [
231 31, 48, 29, 177, 88, 236, 186, 84, 65, 51, 214, 243, 174, 24, 45, 101, 229, 129, 62,
232 132, 45, 174, 183, 65, 89, 73, 107, 177, 77, 90, 164, 251,
233 ];
234
235 let mut alice = ClientHandshake::<Neko128, MlKem512>::new(Some(&psk));
236 let mut bob = ServerHandshake::<Neko256, MlKem512>::new(Some(&psk));
237
238 // BufferSlice acts as our transport across the webz
239 let mut buf = alloc::vec![0u8; 2048];
240 let mut buf = BufferSlice::new(&mut buf);
241
242 let mut rng = rand_core::UnwrapErr(getrandom::SysRng);
243
244 alice.send(&mut rng, &mut buf).unwrap();
245
246 // Pretend to send ek across the webz: client -> server
247 let _bob = bob.respond(&mut rng, &mut buf).unwrap();
248 }
249
250 #[test]
251 #[should_panic]
252 fn handshake_fails_with_kem_security_mismatch() {
253 let psk: [u8; 32] = [
254 31, 48, 29, 177, 88, 236, 186, 84, 65, 51, 214, 243, 174, 24, 45, 101, 229, 129, 62,
255 132, 45, 174, 183, 65, 89, 73, 107, 177, 77, 90, 164, 251,
256 ];
257
258 let mut alice = ClientHandshake::<Neko128, MlKem512>::new(Some(&psk));
259 let mut bob = ServerHandshake::<Neko128, MlKem768>::new(Some(&psk));
260
261 // BufferSlice acts as our transport across the webz
262 let mut buf = alloc::vec![0u8; 2048];
263 let mut buf = BufferSlice::new(&mut buf);
264
265 let mut rng = rand_core::UnwrapErr(getrandom::SysRng);
266
267 alice.send(&mut rng, &mut buf).unwrap();
268
269 // Pretend to send ek across the webz: client -> server
270 let _bob = bob.respond(&mut rng, &mut buf).unwrap();
271 }
272
273 #[test]
274 #[should_panic]
275 fn handshake_fails_with_psk_mismatch() {
276 let psk: [u8; 32] = [
277 31, 48, 29, 177, 88, 236, 186, 84, 65, 51, 214, 243, 174, 24, 45, 101, 229, 129, 62,
278 132, 45, 174, 183, 65, 89, 73, 107, 177, 77, 90, 164, 251,
279 ];
280
281 let mut alice = ClientHandshake::<Neko128, MlKem512>::new(None);
282 let mut bob = ServerHandshake::<Neko128, MlKem512>::new(Some(&psk));
283
284 // BufferSlice acts as our transport across the webz
285 let mut buf = alloc::vec![0u8; 2048];
286 let mut buf = BufferSlice::new(&mut buf);
287
288 let mut rng = rand_core::UnwrapErr(getrandom::SysRng);
289
290 alice.send(&mut rng, &mut buf).unwrap();
291
292 // Pretend to send ek across the webz: client -> server
293 let _bob = bob.respond(&mut rng, &mut buf).unwrap();
294 }
295}