Skip to main content

agent_phone/
noise.rs

1//! Noise_XK_25519_ChaChaPoly_BLAKE2s — direct port of the TS reference.
2
3use 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
65/// X25519 scalar mult via curve25519-dalek's MontgomeryPoint.
66fn 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        // -> e, es
254        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        // Responder: first split key is the initiator→responder one (recv side).
376        let (recv_cs, send_cs) = self.ss.split();
377        HandshakeResult { send_cs, recv_cs }
378    }
379}