morf/
peer.rs

1use chacha20::{cipher::StreamCipher, ChaCha20};
2use chacha20poly1305::{AeadInPlace, ChaCha20Poly1305, ChaChaPoly1305, KeyInit};
3use x25519_dalek::{PublicKey, StaticSecret};
4
5// for early_kdf
6const INITIAL_ENCRYPTION_KEY_INFO: &str = "initial_encryption_key";
7const SERVER_EPHEMERAL_PUBLIC_KEY_MAC_KEY_INFO: &str = "server_ephemeral_public_key_mac_key";
8
9// for session_kdf
10const S2C_STREAM_KEY_INFO: &str = "s2c_stream_key";
11const C2S_STREAM_KEY_INFO: &str = "c2s_stream_key";
12
13pub struct MorfPeer {
14  tx_cipher: Option<ChaCha20Poly1305>,
15  rx_cipher: Option<ChaCha20Poly1305>,
16  client_handshake_state: Option<ClientHandshakeState>,
17  tx_counter: u16,
18  rx_counter: u16,
19}
20
21struct ClientHandshakeState {
22  ephemeral_secret: StaticSecret,
23  server_ephemeral_key_mac_key: [u8; 32],
24}
25
26#[derive(Clone)]
27pub struct UnauthenticatedClientHandshake {
28  pub device_static_public_key_hash: [u8; 16],
29
30  pub server_ephemeral_public_key: PublicKey,
31  pub server_ephemeral_secret_key: StaticSecret,
32  pub client_ephemeral_public_key: PublicKey,
33}
34
35impl UnauthenticatedClientHandshake {
36  pub fn authenticate(self, device_static_public_key: &PublicKey, peer: &mut MorfPeer) {
37    let mut session_key_material = [0u8; 64];
38    session_key_material[0..32].copy_from_slice(
39      self
40        .server_ephemeral_secret_key
41        .diffie_hellman(device_static_public_key)
42        .as_bytes(),
43    );
44    session_key_material[32..64].copy_from_slice(
45      self
46        .server_ephemeral_secret_key
47        .diffie_hellman(&self.client_ephemeral_public_key)
48        .as_bytes(),
49    );
50
51    let tx_stream_key: [u8; 32] = blake3::derive_key(S2C_STREAM_KEY_INFO, &session_key_material);
52    let rx_stream_key: [u8; 32] = blake3::derive_key(C2S_STREAM_KEY_INFO, &session_key_material);
53
54    assert!(peer.tx_cipher.is_none() && peer.rx_cipher.is_none());
55    assert!(peer.client_handshake_state.is_none());
56
57    peer.tx_cipher = Some(ChaChaPoly1305::new(&tx_stream_key.into()));
58    peer.rx_cipher = Some(ChaChaPoly1305::new(&rx_stream_key.into()));
59  }
60}
61
62impl MorfPeer {
63  pub fn client_initiate_handshake(
64    ephemeral_secret: [u8; 32],
65    server_public_key: &PublicKey,
66    device_static_public_key: &PublicKey,
67  ) -> (Self, impl AsMut<[u8]> + AsRef<[u8]>) {
68    let ephemeral_secret = StaticSecret::from(ephemeral_secret);
69    let ephemeral_public = PublicKey::from(&ephemeral_secret);
70    let early_key_material = ephemeral_secret.diffie_hellman(server_public_key);
71
72    let initial_encryption_key: [u8; 32] =
73      blake3::derive_key(INITIAL_ENCRYPTION_KEY_INFO, early_key_material.as_bytes());
74    let mut init_cipher = <ChaCha20 as chacha20::cipher::KeyIvInit>::new(
75      &initial_encryption_key.into(),
76      &[0u8; 12].into(),
77    );
78
79    let device_static_public_key_hash =
80      &mut (<[u8; 32]>::from(blake3::hash(device_static_public_key.as_bytes())))[..16];
81    init_cipher.apply_keystream(device_static_public_key_hash);
82
83    let mut output = [0u8; 49];
84    output[0] = 3;
85    output[1..33].copy_from_slice(ephemeral_public.as_bytes());
86    output[33..49].copy_from_slice(&device_static_public_key_hash);
87
88    let server_ephemeral_key_mac_key = blake3::derive_key(
89      SERVER_EPHEMERAL_PUBLIC_KEY_MAC_KEY_INFO,
90      early_key_material.as_bytes(),
91    );
92
93    (
94      Self {
95        tx_cipher: None,
96        rx_cipher: None,
97        client_handshake_state: Some(ClientHandshakeState {
98          ephemeral_secret,
99          server_ephemeral_key_mac_key,
100        }),
101        tx_counter: 0,
102        rx_counter: 0,
103      },
104      output,
105    )
106  }
107
108  pub fn server_accept_handshake(
109    ephemeral_secret: [u8; 32],
110    static_secret: &StaticSecret,
111    packet: &[u8],
112  ) -> Result<
113    (
114      Self,
115      UnauthenticatedClientHandshake,
116      impl AsMut<[u8]> + AsRef<[u8]>,
117    ),
118    (),
119  > {
120    if packet.len() != 49 || packet[0] != 3 {
121      return Err(());
122    }
123
124    let client_ephemeral_public_key =
125      PublicKey::from(<[u8; 32]>::try_from(&packet[1..33]).unwrap());
126    let mut device_static_public_key_hash: [u8; 16] = packet[33..49].try_into().unwrap();
127
128    let ephemeral_secret = StaticSecret::from(ephemeral_secret);
129    let ephemeral_public = PublicKey::from(&ephemeral_secret);
130
131    let early_key_material = static_secret.diffie_hellman(&client_ephemeral_public_key);
132
133    let mut init_cipher = <ChaCha20 as chacha20::cipher::KeyIvInit>::new(
134      &blake3::derive_key(INITIAL_ENCRYPTION_KEY_INFO, early_key_material.as_bytes()).into(),
135      &[0u8; 12].into(),
136    );
137
138    init_cipher.apply_keystream(&mut device_static_public_key_hash);
139
140    // Derive the MAC key for server ephemeral public key
141    let server_ephemeral_key_mac_key = blake3::derive_key(
142      SERVER_EPHEMERAL_PUBLIC_KEY_MAC_KEY_INFO,
143      early_key_material.as_bytes(),
144    );
145    let mut output = [0u8; 49];
146    output[0] = 1;
147    output[1..33].copy_from_slice(ephemeral_public.as_bytes());
148    output[33..49].copy_from_slice(
149      &blake3::keyed_hash(&server_ephemeral_key_mac_key, ephemeral_public.as_bytes()).as_bytes()
150        [..16],
151    );
152
153    Ok((
154      Self {
155        tx_cipher: None,
156        rx_cipher: None,
157        client_handshake_state: None,
158        tx_counter: 0,
159        rx_counter: 0,
160      },
161      UnauthenticatedClientHandshake {
162        device_static_public_key_hash,
163        server_ephemeral_public_key: ephemeral_public,
164        server_ephemeral_secret_key: ephemeral_secret,
165        client_ephemeral_public_key,
166      },
167      output,
168    ))
169  }
170
171  pub fn client_finalize_handshake(
172    &mut self,
173    device_static_secret_key: &StaticSecret,
174    packet: &[u8],
175  ) -> Result<(), ()> {
176    if packet.len() != 49 || packet[0] != 1 {
177      return Err(());
178    }
179
180    let Some(st) = &self.client_handshake_state else {
181      return Err(());
182    };
183
184    if !constant_time_eq::constant_time_eq(
185      &blake3::keyed_hash(&st.server_ephemeral_key_mac_key, &packet[1..33]).as_bytes()[..16],
186      &packet[33..49],
187    ) {
188      return Err(());
189    }
190
191    // at this point the packet is authenticated
192
193    let server_ephemeral_public_key =
194      PublicKey::from(<[u8; 32]>::try_from(&packet[1..33]).unwrap());
195
196    let mut session_key_material = [0u8; 64];
197    session_key_material[0..32].copy_from_slice(
198      device_static_secret_key
199        .diffie_hellman(&server_ephemeral_public_key)
200        .as_bytes(),
201    );
202    session_key_material[32..64].copy_from_slice(
203      st.ephemeral_secret
204        .diffie_hellman(&server_ephemeral_public_key)
205        .as_bytes(),
206    );
207
208    let tx_stream_key: [u8; 32] = blake3::derive_key(C2S_STREAM_KEY_INFO, &session_key_material);
209    let rx_stream_key: [u8; 32] = blake3::derive_key(S2C_STREAM_KEY_INFO, &session_key_material);
210
211    self.tx_cipher = Some(ChaChaPoly1305::new(&tx_stream_key.into()));
212    self.rx_cipher = Some(ChaChaPoly1305::new(&rx_stream_key.into()));
213    self.client_handshake_state = None;
214
215    Ok(())
216  }
217
218  pub fn unseal<'a>(&mut self, packet: &'a mut [u8]) -> Result<&'a mut [u8], ()> {
219    let Some(rx_cipher) = &self.rx_cipher else {
220      return Err(());
221    };
222    if packet.len() < 19 || packet[0] != 2 {
223      return Err(());
224    }
225
226    let packet_len = packet.len();
227
228    let untrusted_counter = u16::from_be_bytes(packet[1..3].try_into().unwrap());
229    let tag = <[u8; 16]>::try_from(&packet[packet_len - 16..]).unwrap();
230    let content = &mut packet[3..packet_len - 16];
231
232    if untrusted_counter <= self.rx_counter {
233      return Err(());
234    }
235
236    let mut nonce = [0u8; 12];
237    nonce[10..12].copy_from_slice(&untrusted_counter.to_be_bytes());
238
239    rx_cipher
240      .decrypt_in_place_detached(&nonce.into(), &[], content, &tag.into())
241      .map_err(|_| ())?;
242
243    self.rx_counter = untrusted_counter;
244    Ok(content)
245  }
246
247  pub fn seal(&mut self, packet: &mut [u8]) -> Result<([u8; 3], [u8; 16]), ()> {
248    let Some(tx_cipher) = &self.tx_cipher else {
249      return Err(());
250    };
251
252    let next_counter = match self.tx_counter.checked_add(1) {
253      Some(next_counter) => next_counter,
254      None => return Err(()),
255    };
256    self.tx_counter = next_counter;
257
258    let mut nonce = [0u8; 12];
259    nonce[10..12].copy_from_slice(&self.tx_counter.to_be_bytes());
260    let tag = tx_cipher
261      .encrypt_in_place_detached(&nonce.into(), &[], packet)
262      .unwrap();
263
264    let mut prefix = [0u8; 3];
265    prefix[0] = 2;
266    prefix[1..3].copy_from_slice(&self.tx_counter.to_be_bytes());
267    Ok((prefix, tag.into()))
268  }
269}
morf/peer.rs

morf/
peer.rs
