sachy.dev
Repo of no-std crates for my personal embedded projects
1#![no_std]
2
3use core::ops::BitXor;
4
5use chacha20poly1305::{AeadInOut, ChaCha20Poly1305, KeyInit, aead};
6use dhkem::{
7 Encapsulate, Expander, Kem, X25519DecapsulationKey, X25519EncapsulationKey, X25519Kem,
8 kem::{Ciphertext, Decapsulate, Key, KeyExport, SharedKey, TryKeyInit},
9};
10use elliptic_curve::subtle::ConstantTimeEq;
11
12/// Error type.
13///
14/// This type is deliberately opaque as to avoid potential side-channel
15/// leakage (e.g. padding oracle).
16#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
17pub struct ProtoError;
18
19impl core::fmt::Display for ProtoError {
20 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
21 f.write_str("ProtoError")
22 }
23}
24
25impl core::error::Error for ProtoError {}
26
27impl From<chacha20poly1305::Error> for ProtoError {
28 fn from(_value: chacha20poly1305::Error) -> Self {
29 Self
30 }
31}
32
33pub struct ClientHandshake(X25519DecapsulationKey);
34
35pub struct EncapsulatedPublicKey(X25519EncapsulationKey);
36
37/// The role of the participant, whether sending/receiving during handshake,
38/// and then whether sending/receiving during communication.
39#[derive(Debug, PartialEq, Eq, Clone, Copy)]
40pub enum Role {
41 /// Participant SENDS data
42 Sender,
43 /// Participant RECEIVES data
44 Receiver,
45}
46
47impl From<Role> for u8 {
48 fn from(value: Role) -> Self {
49 match value {
50 Role::Sender => 0,
51 Role::Receiver => 1,
52 }
53 }
54}
55
56impl BitXor for Role {
57 type Output = u8;
58
59 fn bitxor(self, rhs: Self) -> u8 {
60 u8::from(self) ^ u8::from(rhs)
61 }
62}
63
64impl EncapsulatedPublicKey {
65 pub fn serialize(&self) -> Key<X25519EncapsulationKey> {
66 self.0.to_bytes()
67 }
68
69 pub fn deserialize(buf: &[u8]) -> Result<Self, ProtoError> {
70 Ok(Self(
71 X25519EncapsulationKey::new_from_slice(buf).map_err(|_| ProtoError)?,
72 ))
73 }
74
75 pub fn encapsulate(&self) -> (Ciphertext<X25519Kem>, SharedKey<X25519Kem>) {
76 self.0.encapsulate()
77 }
78}
79
80impl ClientHandshake {
81 pub fn send() -> (EncapsulatedPublicKey, Self) {
82 let (decap, encap) = X25519Kem::generate_keypair();
83
84 (EncapsulatedPublicKey(encap), Self(decap))
85 }
86
87 pub fn finish(self, ciphertext: &[u8], psk: &[u8; 32]) -> Result<TransportState, ProtoError> {
88 let shared = self
89 .0
90 .decapsulate_slice(ciphertext)
91 .map_err(|_| ProtoError)?;
92
93 TransportState::init(psk, shared, Role::Sender)
94 }
95}
96
97pub struct ServerHandshake(SharedKey<X25519Kem>);
98
99impl ServerHandshake {
100 pub fn receive(buf: &[u8]) -> Result<(Ciphertext<X25519Kem>, Self), ProtoError> {
101 let encap = EncapsulatedPublicKey::deserialize(buf)?;
102
103 let (ciphertext, sk) = encap.encapsulate();
104
105 Ok((ciphertext, Self(sk)))
106 }
107
108 pub fn finish(self, psk: &[u8; 32]) -> Result<TransportState, ProtoError> {
109 TransportState::init(psk, self.0, Role::Receiver)
110 }
111}
112
113pub struct SendingState<'a> {
114 transport: &'a TransportState,
115 counter: u64,
116}
117
118impl SendingState<'_> {
119 pub fn encrypt(
120 &mut self,
121 msg: &mut dyn aead::Buffer,
122 associated_data: &[u8],
123 ) -> Result<(), ProtoError> {
124 if self.counter.ct_eq(&u64::MAX).into() {
125 return Err(ProtoError);
126 }
127
128 self.transport.aead.encrypt_in_place(
129 &self
130 .transport
131 .mix_nonce(&self.counter.to_be_bytes(), Role::Sender),
132 associated_data,
133 msg,
134 )?;
135
136 self.counter = self.counter.wrapping_add(1);
137
138 Ok(())
139 }
140}
141
142pub struct ReceivingState<'a> {
143 transport: &'a TransportState,
144 counter: u64,
145}
146
147impl ReceivingState<'_> {
148 pub fn decrypt(
149 &mut self,
150 msg: &mut dyn aead::Buffer,
151 associated_data: &[u8],
152 ) -> Result<(), ProtoError> {
153 if self.counter.ct_eq(&u64::MAX).into() {
154 return Err(ProtoError);
155 }
156
157 self.transport.aead.decrypt_in_place(
158 &self
159 .transport
160 .mix_nonce(&self.counter.to_be_bytes(), Role::Receiver),
161 associated_data,
162 msg,
163 )?;
164
165 self.counter = self.counter.wrapping_add(1);
166
167 Ok(())
168 }
169}
170
171#[repr(align(4))]
172pub struct TransportState {
173 aead: ChaCha20Poly1305,
174 client: aead::Nonce<ChaCha20Poly1305>,
175 server: aead::Nonce<ChaCha20Poly1305>,
176 role: Role,
177}
178
179impl TransportState {
180 pub fn init(
181 psk: &[u8; 32],
182 shared: SharedKey<X25519Kem>,
183 role: Role,
184 ) -> Result<Self, ProtoError> {
185 let kdf = Expander::<sha2::Sha256>::new_labeled_hpke(psk, b"Sachy-Crypto", &shared)
186 .map_err(|_| ProtoError)?;
187
188 let mut key = [0u8; 32];
189 let mut client = aead::Nonce::<ChaCha20Poly1305>::default();
190 let mut server = aead::Nonce::<ChaCha20Poly1305>::default();
191
192 kdf.expand(b"SecretKey012", &mut key)
193 .map_err(|_| ProtoError)?;
194 kdf.expand(b"NonceClient*", &mut client)
195 .map_err(|_| ProtoError)?;
196 kdf.expand(b"NonceServer#", &mut server)
197 .map_err(|_| ProtoError)?;
198
199 Ok(Self {
200 aead: ChaCha20Poly1305::new(&key.into()),
201 client,
202 server,
203 role,
204 })
205 }
206
207 pub fn split(&self) -> (SendingState<'_>, ReceivingState<'_>) {
208 (
209 SendingState {
210 transport: self,
211 counter: 0,
212 },
213 ReceivingState {
214 transport: self,
215 counter: 0,
216 },
217 )
218 }
219
220 /// Selects which nonce to use for encrypting/decrypting, which matters for
221 /// ensuring the same nonce is used only for one direction of communication.
222 fn select_nonce_context(&self, send: Role) -> &aead::Nonce<ChaCha20Poly1305> {
223 let context_select = self.role ^ send;
224
225 // Handshake ROLE XOR Transport ROLE selects either one or other nonce context,
226 // (0) for first context, (1) for second context
227 // Sending: Client ^ Sender = 0 (select first/client context)
228 // Receiving: Server ^ Receiver = 0 (select first/client context)
229 // Sending: Server ^ Sender = 1 (select second/server context)
230 // Receiving: Client ^ Receiver = 1 (select second/server context)
231 if context_select.ct_eq(&0).into() {
232 &self.client
233 } else {
234 &self.server
235 }
236 }
237
238 fn mix_nonce(&self, position: &[u8; 8], send: Role) -> aead::Nonce<ChaCha20Poly1305> {
239 let mut trump = aead::Nonce::<ChaCha20Poly1305>::default();
240
241 let epstein = self.select_nonce_context(send);
242
243 let index = trump.len() - position.len();
244
245 // Copy position bytes into BE format onto derived nonce
246 trump[index..].copy_from_slice(position);
247
248 // XOR the base nonce onto the derived nonce bytes
249 trump
250 .iter_mut()
251 .zip(epstein)
252 .for_each(|(trump, epstein)| *trump ^= *epstein);
253
254 trump
255 }
256}
257
258#[derive(Debug)]
259pub struct BufferSlice<'a> {
260 slice: &'a mut [u8],
261 end: usize,
262}
263
264impl<'a> BufferSlice<'a> {
265 pub fn new(slice: &'a mut [u8]) -> Self {
266 Self {
267 end: slice.len(),
268 slice,
269 }
270 }
271
272 pub fn reset(&mut self) {
273 self.end = self.slice.len();
274 }
275}
276
277impl AsRef<[u8]> for BufferSlice<'_> {
278 fn as_ref(&self) -> &[u8] {
279 &self.slice[..self.end]
280 }
281}
282
283impl AsMut<[u8]> for BufferSlice<'_> {
284 fn as_mut(&mut self) -> &mut [u8] {
285 &mut self.slice[..self.end]
286 }
287}
288
289impl aead::Buffer for BufferSlice<'_> {
290 fn extend_from_slice(&mut self, other: &[u8]) -> aead::Result<()> {
291 let index = self.end + other.len();
292
293 if index > self.slice.len() {
294 return Err(aead::Error);
295 }
296
297 self.slice[self.end..index].copy_from_slice(other);
298
299 self.end = index;
300
301 Ok(())
302 }
303
304 fn truncate(&mut self, len: usize) {
305 self.end = len;
306 }
307}
308
309#[cfg(test)]
310mod tests {
311 use alloc::vec;
312 use chacha20poly1305::aead::Buffer;
313 use dhkem::Generate;
314 use elliptic_curve::array::Array;
315
316 extern crate alloc;
317
318 use super::*;
319
320 #[test]
321 fn buffer_slice_works() {
322 let mut buf = vec![0u8; 128];
323
324 let mut buf_slice = BufferSlice::new(&mut buf);
325
326 assert_eq!(buf_slice.len(), 128);
327 assert_eq!(buf_slice.extend_from_slice(&[0, 0, 0]), Err(aead::Error));
328
329 buf_slice.truncate(64);
330
331 assert_eq!(buf_slice.extend_from_slice(&[0, 0, 0, 0, 0, 0]), Ok(()));
332 assert_eq!(buf_slice.len(), 70);
333
334 buf_slice.reset();
335
336 assert_eq!(buf_slice.len(), 128);
337 }
338
339 #[test]
340 fn handshake_protocol_works() -> Result<(), ProtoError> {
341 let psk: [u8; 32] = [
342 31, 48, 29, 177, 88, 236, 186, 84, 65, 51, 214, 243, 174, 24, 45, 101, 229, 129, 62,
343 132, 45, 174, 183, 65, 89, 73, 107, 177, 77, 90, 164, 251,
344 ];
345
346 let (ek, client) = ClientHandshake::send();
347
348 // Pretend to send ek across the webz: client -> server
349 let (ciphertext, server) = ServerHandshake::receive(&ek.serialize())?;
350
351 // Pretend to send ciphertext across the webz: server -> client
352 let alice = client.finish(&ciphertext, &psk)?;
353 let bob = server.finish(&psk)?;
354
355 let nonce = aead::Nonce::<ChaCha20Poly1305>::generate();
356
357 let mut buffer1 = vec![0u8; 64];
358 let mut buffer2 = vec![0u8; 64];
359
360 // Using the same nonce to check that the internal AEAD states match. Normally, client/server
361 // would work with unique derived nonces, because nonce reuse is BAD
362 alice.aead.encrypt_in_place(&nonce, &[], &mut buffer1)?;
363 bob.aead.encrypt_in_place(&nonce, &[], &mut buffer2)?;
364
365 // If the nonces match, then we can assume the rest of the internal state is the same too
366 // so the outputs should match each other
367 assert_eq!(&buffer1, &buffer2);
368
369 // Both Transports have derived base nonces for each context.
370 // Client context nonces will not match Server context nonces.
371 assert_eq!(alice.client, bob.client);
372 assert_eq!(alice.server, bob.server);
373 assert_ne!(alice.client, alice.server);
374 assert_ne!(bob.client, bob.server);
375
376 Ok(())
377 }
378
379 #[test]
380 fn two_way_transport_sync_works() -> Result<(), ProtoError> {
381 let shared_secret = [
382 0x80, 0x81, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8a, 0x8b, 0x8c, 0x8d,
383 0x8e, 0x8f, 0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9a, 0x9b,
384 0x9c, 0x9d, 0x9e, 0x9f,
385 ];
386
387 let psk: [u8; 32] = [
388 31, 48, 29, 177, 88, 236, 186, 84, 65, 51, 214, 243, 174, 24, 45, 101, 229, 129, 62,
389 132, 45, 174, 183, 65, 89, 73, 107, 177, 77, 90, 164, 251,
390 ];
391
392 let alice = TransportState::init(&psk, Array(shared_secret), Role::Sender)?;
393 let bob = TransportState::init(&psk, Array(shared_secret), Role::Receiver)?;
394
395 let (mut alice_send, mut alice_recv) = alice.split();
396 let (mut bob_send, mut bob_recv) = bob.split();
397
398 let orig = b"Test Message, Please ignore.";
399
400 let ad = b"random";
401
402 let mut msg = orig.to_vec();
403
404 // a -> b
405 alice_send.encrypt(&mut msg, ad)?;
406
407 assert_ne!(orig.as_slice(), msg.as_slice());
408 let ct1 = msg.clone();
409
410 bob_recv.decrypt(&mut msg, ad)?;
411
412 // a -> b
413 alice_send.encrypt(&mut msg, b"")?;
414
415 assert_ne!(msg.as_slice(), ct1.as_slice());
416 let ct2 = msg.clone();
417
418 bob_recv.decrypt(&mut msg, b"")?;
419
420 // b -> a
421 bob_send.encrypt(&mut msg, ad)?;
422
423 // None of the ciphertexts should match each other
424 assert_ne!(msg.as_slice(), ct1.as_slice());
425 assert_ne!(msg.as_slice(), ct2.as_slice());
426 assert_ne!(ct1.as_slice(), ct2.as_slice());
427
428 alice_recv.decrypt(&mut msg, ad)?;
429
430 assert_eq!(orig.as_slice(), msg.as_slice());
431
432 // Counters are tracked from sender to receiver
433 assert_eq!(alice_send.counter, bob_recv.counter);
434 assert_eq!(bob_send.counter, alice_recv.counter);
435
436 // Counters are not linked on the same side
437 assert_ne!(alice_send.counter, alice_recv.counter);
438 assert_ne!(bob_send.counter, bob_recv.counter);
439
440 Ok(())
441 }
442}