1use crate::error::{Error, Result};
4use blake2::{Blake2s256, Digest};
5use chacha20poly1305::{
6 aead::{Aead, KeyInit, Payload},
7 ChaCha20Poly1305, Key, Nonce,
8};
9use curve25519_dalek::montgomery::MontgomeryPoint;
10use curve25519_dalek::scalar::clamp_integer;
11use rand::RngCore;
12
13pub const PROTOCOL: &[u8] = b"Noise_XK_25519_ChaChaPoly_BLAKE2s";
14pub const HASHLEN: usize = 32;
15pub const BLOCKLEN: usize = 64;
16
17fn blake2s(data: &[u8]) -> [u8; 32] {
18 let mut h = Blake2s256::new();
19 h.update(data);
20 h.finalize().into()
21}
22
23pub fn hmac_blake2s(key: &[u8], data: &[u8]) -> [u8; 32] {
24 let k_bytes;
25 let k: &[u8] = if key.len() > BLOCKLEN {
26 k_bytes = blake2s(key);
27 &k_bytes
28 } else {
29 key
30 };
31 let mut padded = [0u8; BLOCKLEN];
32 padded[..k.len()].copy_from_slice(k);
33 let mut ipad = [0u8; BLOCKLEN];
34 let mut opad = [0u8; BLOCKLEN];
35 for i in 0..BLOCKLEN {
36 ipad[i] = padded[i] ^ 0x36;
37 opad[i] = padded[i] ^ 0x5C;
38 }
39 let mut inner = Vec::with_capacity(BLOCKLEN + data.len());
40 inner.extend_from_slice(&ipad);
41 inner.extend_from_slice(data);
42 let inner_h = blake2s(&inner);
43 let mut outer = Vec::with_capacity(BLOCKLEN + HASHLEN);
44 outer.extend_from_slice(&opad);
45 outer.extend_from_slice(&inner_h);
46 blake2s(&outer)
47}
48
49pub fn hkdf2(ck: &[u8], ikm: &[u8]) -> ([u8; 32], [u8; 32]) {
50 let t0 = hmac_blake2s(ck, ikm);
51 let t1 = hmac_blake2s(&t0, &[0x01]);
52 let mut t2_input = Vec::with_capacity(33);
53 t2_input.extend_from_slice(&t1);
54 t2_input.push(0x02);
55 let t2 = hmac_blake2s(&t0, &t2_input);
56 (t1, t2)
57}
58
59fn nonce_bytes(n: u64) -> [u8; 12] {
60 let mut out = [0u8; 12];
61 out[4..].copy_from_slice(&n.to_le_bytes());
62 out
63}
64
65fn x25519(scalar: &[u8; 32], u: &[u8; 32]) -> [u8; 32] {
67 let clamped = clamp_integer(*scalar);
68 let scalar_inner = curve25519_dalek::scalar::Scalar::from_bytes_mod_order(clamped);
69 let point = MontgomeryPoint(*u);
70 (point * scalar_inner).to_bytes()
71}
72
73pub fn x25519_public_from_private(priv_: &[u8; 32]) -> [u8; 32] {
74 let mut base = [0u8; 32];
75 base[0] = 9;
76 x25519(priv_, &base)
77}
78
79pub fn x25519_random_private() -> [u8; 32] {
80 let mut k = [0u8; 32];
81 rand::thread_rng().fill_bytes(&mut k);
82 k
83}
84
85#[derive(Clone)]
86pub struct CipherState {
87 pub k: Option<[u8; 32]>,
88 pub n: u64,
89}
90
91impl CipherState {
92 pub fn new(k: Option<[u8; 32]>) -> Self {
93 Self { k, n: 0 }
94 }
95
96 pub fn encrypt_with_ad(&mut self, ad: &[u8], plaintext: &[u8]) -> Result<Vec<u8>> {
97 let key = match self.k {
98 Some(k) => k,
99 None => return Ok(plaintext.to_vec()),
100 };
101 let aead = ChaCha20Poly1305::new(Key::from_slice(&key));
102 let nonce = nonce_bytes(self.n);
103 let ct = aead
104 .encrypt(
105 Nonce::from_slice(&nonce),
106 Payload {
107 msg: plaintext,
108 aad: ad,
109 },
110 )
111 .map_err(|e| Error::Noise(format!("encrypt: {e}")))?;
112 self.n += 1;
113 Ok(ct)
114 }
115
116 pub fn decrypt_with_ad(&mut self, ad: &[u8], ciphertext: &[u8]) -> Result<Vec<u8>> {
117 let key = match self.k {
118 Some(k) => k,
119 None => return Ok(ciphertext.to_vec()),
120 };
121 let aead = ChaCha20Poly1305::new(Key::from_slice(&key));
122 let nonce = nonce_bytes(self.n);
123 let pt = aead
124 .decrypt(
125 Nonce::from_slice(&nonce),
126 Payload {
127 msg: ciphertext,
128 aad: ad,
129 },
130 )
131 .map_err(|e| Error::Noise(format!("decrypt: {e}")))?;
132 self.n += 1;
133 Ok(pt)
134 }
135}
136
137pub struct SymmetricState {
138 pub ck: [u8; 32],
139 pub h: [u8; 32],
140 pub cs: CipherState,
141}
142
143impl SymmetricState {
144 pub fn new() -> Self {
145 let mut h = [0u8; HASHLEN];
146 if PROTOCOL.len() <= HASHLEN {
147 h[..PROTOCOL.len()].copy_from_slice(PROTOCOL);
148 } else {
149 h = blake2s(PROTOCOL);
150 }
151 let ck = h;
152 Self {
153 ck,
154 h,
155 cs: CipherState::new(None),
156 }
157 }
158
159 pub fn mix_hash(&mut self, data: &[u8]) {
160 let mut buf = Vec::with_capacity(HASHLEN + data.len());
161 buf.extend_from_slice(&self.h);
162 buf.extend_from_slice(data);
163 self.h = blake2s(&buf);
164 }
165
166 pub fn mix_key(&mut self, ikm: &[u8]) {
167 let (new_ck, temp_k) = hkdf2(&self.ck, ikm);
168 self.ck = new_ck;
169 self.cs = CipherState::new(Some(temp_k));
170 }
171
172 pub fn encrypt_and_hash(&mut self, plaintext: &[u8]) -> Result<Vec<u8>> {
173 let ct = self.cs.encrypt_with_ad(&self.h, plaintext)?;
174 self.mix_hash(&ct);
175 Ok(ct)
176 }
177
178 pub fn decrypt_and_hash(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>> {
179 let pt = self.cs.decrypt_with_ad(&self.h, ciphertext)?;
180 self.mix_hash(ciphertext);
181 Ok(pt)
182 }
183
184 pub fn split(self) -> (CipherState, CipherState) {
185 let (k1, k2) = hkdf2(&self.ck, &[]);
186 (CipherState::new(Some(k1)), CipherState::new(Some(k2)))
187 }
188}
189
190impl Default for SymmetricState {
191 fn default() -> Self {
192 Self::new()
193 }
194}
195
196pub struct HandshakeResult {
197 pub send_cs: CipherState,
198 pub recv_cs: CipherState,
199}
200
201impl HandshakeResult {
202 pub fn send(&mut self, plaintext: &[u8]) -> Result<Vec<u8>> {
203 self.send_cs.encrypt_with_ad(&[], plaintext)
204 }
205 pub fn recv(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>> {
206 self.recv_cs.decrypt_with_ad(&[], ciphertext)
207 }
208}
209
210pub fn build_prologue(initiator_did: &str, responder_did: &str) -> Vec<u8> {
211 let prefix = b"agent-phone/1";
212 let init = initiator_did.as_bytes();
213 let resp = responder_did.as_bytes();
214 let mut out = Vec::with_capacity(prefix.len() + 2 + init.len() + 2 + resp.len());
215 out.extend_from_slice(prefix);
216 out.extend_from_slice(&(init.len() as u16).to_be_bytes());
217 out.extend_from_slice(init);
218 out.extend_from_slice(&(resp.len() as u16).to_be_bytes());
219 out.extend_from_slice(resp);
220 out
221}
222
223pub struct InitiatorHandshake {
224 ss: SymmetricState,
225 static_priv: [u8; 32],
226 static_pub: [u8; 32],
227 responder_static_pub: [u8; 32],
228 e_priv: Option<[u8; 32]>,
229 re_pub: Option<[u8; 32]>,
230}
231
232impl InitiatorHandshake {
233 pub fn new(
234 prologue: &[u8],
235 static_priv: [u8; 32],
236 static_pub: [u8; 32],
237 responder_static_pub: [u8; 32],
238 ) -> Self {
239 let mut ss = SymmetricState::new();
240 ss.mix_hash(prologue);
241 ss.mix_hash(&responder_static_pub);
242 Self {
243 ss,
244 static_priv,
245 static_pub,
246 responder_static_pub,
247 e_priv: None,
248 re_pub: None,
249 }
250 }
251
252 pub fn write_message_1(&mut self) -> Result<Vec<u8>> {
253 let e_priv = x25519_random_private();
255 let e_pub = x25519_public_from_private(&e_priv);
256 self.ss.mix_hash(&e_pub);
257 self.ss
258 .mix_key(&x25519(&e_priv, &self.responder_static_pub));
259 let enc = self.ss.encrypt_and_hash(&[])?;
260 self.e_priv = Some(e_priv);
261 let mut out = Vec::with_capacity(32 + enc.len());
262 out.extend_from_slice(&e_pub);
263 out.extend_from_slice(&enc);
264 Ok(out)
265 }
266
267 pub fn read_message_2(&mut self, msg: &[u8]) -> Result<()> {
268 let e_priv = self
269 .e_priv
270 .ok_or_else(|| Error::Noise("write_message_1 must run first".into()))?;
271 if msg.len() < 32 {
272 return Err(Error::Noise("message 2 too short".into()));
273 }
274 let mut re_pub = [0u8; 32];
275 re_pub.copy_from_slice(&msg[..32]);
276 self.ss.mix_hash(&re_pub);
277 self.ss.mix_key(&x25519(&e_priv, &re_pub));
278 self.ss.decrypt_and_hash(&msg[32..])?;
279 self.re_pub = Some(re_pub);
280 Ok(())
281 }
282
283 pub fn write_message_3(&mut self) -> Result<Vec<u8>> {
284 let re_pub = self
285 .re_pub
286 .ok_or_else(|| Error::Noise("read_message_2 must run first".into()))?;
287 let enc_s = self.ss.encrypt_and_hash(&self.static_pub)?;
288 self.ss.mix_key(&x25519(&self.static_priv, &re_pub));
289 let enc_payload = self.ss.encrypt_and_hash(&[])?;
290 let mut out = Vec::with_capacity(enc_s.len() + enc_payload.len());
291 out.extend_from_slice(&enc_s);
292 out.extend_from_slice(&enc_payload);
293 Ok(out)
294 }
295
296 pub fn finish(self) -> HandshakeResult {
297 let (send_cs, recv_cs) = self.ss.split();
298 HandshakeResult { send_cs, recv_cs }
299 }
300}
301
302pub struct ResponderHandshake {
303 ss: SymmetricState,
304 static_priv: [u8; 32],
305 #[allow(dead_code)]
306 static_pub: [u8; 32],
307 e_priv: Option<[u8; 32]>,
308 re_init_pub: Option<[u8; 32]>,
309}
310
311impl ResponderHandshake {
312 pub fn new(prologue: &[u8], static_priv: [u8; 32], static_pub: [u8; 32]) -> Self {
313 let mut ss = SymmetricState::new();
314 ss.mix_hash(prologue);
315 ss.mix_hash(&static_pub);
316 Self {
317 ss,
318 static_priv,
319 static_pub,
320 e_priv: None,
321 re_init_pub: None,
322 }
323 }
324
325 pub fn read_message_1(&mut self, msg: &[u8]) -> Result<()> {
326 if msg.len() < 32 {
327 return Err(Error::Noise("message 1 too short".into()));
328 }
329 let mut re_init = [0u8; 32];
330 re_init.copy_from_slice(&msg[..32]);
331 self.ss.mix_hash(&re_init);
332 self.ss.mix_key(&x25519(&self.static_priv, &re_init));
333 self.ss.decrypt_and_hash(&msg[32..])?;
334 self.re_init_pub = Some(re_init);
335 Ok(())
336 }
337
338 pub fn write_message_2(&mut self) -> Result<Vec<u8>> {
339 let re_init = self
340 .re_init_pub
341 .ok_or_else(|| Error::Noise("read_message_1 must run first".into()))?;
342 let e_priv = x25519_random_private();
343 let e_pub = x25519_public_from_private(&e_priv);
344 self.ss.mix_hash(&e_pub);
345 self.ss.mix_key(&x25519(&e_priv, &re_init));
346 let enc = self.ss.encrypt_and_hash(&[])?;
347 self.e_priv = Some(e_priv);
348 let mut out = Vec::with_capacity(32 + enc.len());
349 out.extend_from_slice(&e_pub);
350 out.extend_from_slice(&enc);
351 Ok(out)
352 }
353
354 pub fn read_message_3(&mut self, msg: &[u8]) -> Result<()> {
355 let e_priv = self
356 .e_priv
357 .ok_or_else(|| Error::Noise("write_message_2 must run first".into()))?;
358 if msg.len() < 48 {
359 return Err(Error::Noise("message 3 too short".into()));
360 }
361 let enc_s = &msg[..48];
362 let rest = &msg[48..];
363 let ris_bytes = self.ss.decrypt_and_hash(enc_s)?;
364 if ris_bytes.len() != 32 {
365 return Err(Error::Noise("decrypted static key wrong size".into()));
366 }
367 let mut ris = [0u8; 32];
368 ris.copy_from_slice(&ris_bytes);
369 self.ss.mix_key(&x25519(&e_priv, &ris));
370 self.ss.decrypt_and_hash(rest)?;
371 Ok(())
372 }
373
374 pub fn finish(self) -> HandshakeResult {
375 let (recv_cs, send_cs) = self.ss.split();
377 HandshakeResult { send_cs, recv_cs }
378 }
379}