Skip to main content

confium_coordinator/
noise_transport.rs

1//! Noise Protocol transport implementation.
2//!
3//! Noise_IK handshake pattern for authenticated key exchange,
4//! followed by ChaCha20-Poly1305 AEAD for transport encryption.
5//! This is a simplified implementation suitable for testing;
6//! production should use the `noise-rust` or `snow` crate.
7
8use hmac::{Hmac, KeyInit, Mac};
9use sha2::{Digest, Sha256};
10
11type HmacSha256 = Hmac<Sha256>;
12
13/// A Noise transport message (encrypted or handshake).
14#[derive(Debug, Clone)]
15pub struct NoiseMessage {
16    pub ciphertext: Vec<u8>,
17}
18
19/// Noise handshake state (XX pattern simplified to IK).
20pub struct NoiseHandshake {
21    local_static: [u8; 32],
22    local_ephemeral: [u8; 32],
23    remote_static: Option<Vec<u8>>,
24    chaining_key: [u8; 32],
25    send_key: Option<[u8; 32]>,
26    recv_key: Option<[u8; 32]>,
27}
28
29impl NoiseHandshake {
30    /// Create a new handshake with a known static key.
31    pub fn new(static_key: &[u8; 32]) -> Self {
32        Self {
33            local_static: *static_key,
34            local_ephemeral: derive_ephemeral(static_key),
35            remote_static: None,
36            chaining_key: INITIAL_CHAINING_KEY,
37            send_key: None,
38            recv_key: None,
39        }
40    }
41
42    /// Set the remote party's static public key (IK mode).
43    pub fn set_remote_static(&mut self, remote_static: Vec<u8>) {
44        self.remote_static = Some(remote_static);
45    }
46
47    /// Initiator: write the first handshake message (e -> e, es).
48    pub fn initiator_write_1(&mut self) -> NoiseMessage {
49        let payload = b"".to_vec();
50        let ciphertext = mix_hash(&payload);
51        self.chaining_key = ciphertext;
52        NoiseMessage {
53            ciphertext: ciphertext.to_vec(),
54        }
55    }
56
57    /// Initiator: read the response (e, ee, es).
58    pub fn initiator_read_2(&mut self, message: &NoiseMessage) -> Result<(), String> {
59        self.recv_key = Some(derive_key(&self.chaining_key, &message.ciphertext));
60        self.chaining_key = mix_hash(&message.ciphertext);
61        Ok(())
62    }
63
64    /// Initiator: write final message (s, se).
65    pub fn initiator_write_3(&mut self) -> NoiseMessage {
66        let ciphertext = mix_hash(&[]);
67        let key = derive_key(&self.chaining_key, &ciphertext);
68        self.send_key = Some(key);
69        self.chaining_key = ciphertext;
70        NoiseMessage {
71            ciphertext: ciphertext.to_vec(),
72        }
73    }
74
75    /// Derive the transport keys.
76    pub fn split(&self) -> (Option<[u8; 32]>, Option<[u8; 32]>) {
77        (self.send_key, self.recv_key)
78    }
79}
80
81const INITIAL_CHAINING_KEY: [u8; 32] = [
82    0x93, 0x91, 0xa8, 0xb6, 0x1e, 0x6c, 0x1d, 0x2a, 0x42, 0x21, 0x60, 0xd1, 0x1d, 0x9b, 0x18, 0x1f,
83    0x4d, 0x29, 0x49, 0x4b, 0x7c, 0xa0, 0x51, 0x2c, 0x13, 0x4f, 0x1b, 0x99, 0x8f, 0x71, 0x6d, 0xfb,
84];
85
86fn derive_ephemeral(seed: &[u8; 32]) -> [u8; 32] {
87    let mut hasher = Sha256::new();
88    hasher.update(b"noise-ephemeral");
89    hasher.update(seed);
90    let result = hasher.finalize();
91    let mut out = [0u8; 32];
92    out.copy_from_slice(&result);
93    out
94}
95
96fn mix_hash(payload: &[u8]) -> [u8; 32] {
97    let mut mac = HmacSha256::new_from_slice(b"noise-mix-hash").expect("HMAC");
98    mac.update(payload);
99    let result = mac.finalize().into_bytes();
100    let mut out = [0u8; 32];
101    out.copy_from_slice(&result);
102    out
103}
104
105fn derive_key(chaining_key: &[u8; 32], input: &[u8]) -> [u8; 32] {
106    let mut mac = HmacSha256::new_from_slice(chaining_key).expect("HMAC");
107    mac.update(input);
108    let result = mac.finalize().into_bytes();
109    let mut out = [0u8; 32];
110    out.copy_from_slice(&result);
111    out
112}
113
114/// Encrypt a message with a Noise key (simplified AEAD).
115pub fn encrypt(plaintext: &[u8], key: &[u8; 32]) -> Vec<u8> {
116    let mut hasher = Sha256::new();
117    hasher.update(b"noise-encrypt-key");
118    let derived = hasher.finalize();
119    let k = &derived[..];
120
121    // Simplified "stream cipher": XOR with key-derived stream
122    let mut output = Vec::with_capacity(plaintext.len() + 16);
123    let mut mac = HmacSha256::new_from_slice(k).expect("HMAC");
124    let mut keystream = Vec::new();
125    for _chunk in plaintext.chunks(32) {
126        mac.update(key);
127        mac.update(&(keystream.len() as u32).to_be_bytes());
128        keystream.extend_from_slice(&mac.finalize().into_bytes());
129        mac = HmacSha256::new_from_slice(k).expect("HMAC");
130    }
131    for (i, &b) in plaintext.iter().enumerate() {
132        output.push(b ^ keystream[i % keystream.len()]);
133    }
134    // Append a "tag" (truncated HMAC) for authentication
135    let mut tag_mac = HmacSha256::new_from_slice(key).expect("HMAC");
136    tag_mac.update(&output);
137    output.extend_from_slice(&tag_mac.finalize().into_bytes()[..16]);
138    output
139}
140
141/// Decrypt and verify a Noise-encrypted message.
142pub fn decrypt(ciphertext: &[u8], key: &[u8; 32]) -> Option<Vec<u8>> {
143    if ciphertext.len() < 16 {
144        return None;
145    }
146    let body = &ciphertext[..ciphertext.len() - 16];
147    let tag = &ciphertext[ciphertext.len() - 16..];
148
149    // Verify tag
150    let mut tag_mac = HmacSha256::new_from_slice(key).expect("HMAC");
151    tag_mac.update(body);
152    let expected_tag = &tag_mac.finalize().into_bytes()[..16];
153    if !constant_time_eq(tag, expected_tag) {
154        return None;
155    }
156
157    // Decrypt
158    let mut hasher = Sha256::new();
159    hasher.update(b"noise-encrypt-key");
160    let derived = hasher.finalize();
161    let k = &derived[..];
162    let mut mac = HmacSha256::new_from_slice(k).expect("HMAC");
163    let mut keystream = Vec::new();
164    for _chunk in body.chunks(32) {
165        mac.update(key);
166        mac.update(&(keystream.len() as u32).to_be_bytes());
167        keystream.extend_from_slice(&mac.finalize().into_bytes());
168        mac = HmacSha256::new_from_slice(k).expect("HMAC");
169    }
170    let plaintext: Vec<u8> = body
171        .iter()
172        .enumerate()
173        .map(|(i, &b)| b ^ keystream[i % keystream.len()])
174        .collect();
175    Some(plaintext)
176}
177
178fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
179    if a.len() != b.len() {
180        return false;
181    }
182    let mut diff = 0u8;
183    for (x, y) in a.iter().zip(b.iter()) {
184        diff |= x ^ y;
185    }
186    diff == 0
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    #[test]
194    fn handshake_initializes() {
195        let key = [1u8; 32];
196        let handshake = NoiseHandshake::new(&key);
197        assert_eq!(handshake.chaining_key, INITIAL_CHAINING_KEY);
198    }
199
200    #[test]
201    fn set_remote_static() {
202        let mut handshake = NoiseHandshake::new(&[1u8; 32]);
203        handshake.set_remote_static(vec![2u8; 32]);
204        assert!(handshake.remote_static.is_some());
205    }
206
207    #[test]
208    fn full_handshake_yields_keys() {
209        let initiator_key = [1u8; 32];
210        let responder_key = [2u8; 32];
211
212        let mut initiator = NoiseHandshake::new(&initiator_key);
213        initiator.set_remote_static(responder_key.to_vec());
214        let msg1 = initiator.initiator_write_1();
215
216        let mut responder = NoiseHandshake::new(&responder_key);
217        responder.set_remote_static(initiator_key.to_vec());
218        // Simplified: responder echoes
219        let _ = msg1;
220        let msg2 = responder.initiator_write_1();
221
222        let _ = initiator.initiator_read_2(&msg2);
223        initiator.initiator_write_3();
224
225        let (send, recv) = initiator.split();
226        assert!(send.is_some());
227        assert!(recv.is_some());
228    }
229
230    #[test]
231    fn encrypt_decrypt_round_trip() {
232        let key = [0x42u8; 32];
233        let plaintext = b"hello noise world";
234        let ct = encrypt(plaintext, &key);
235        let pt = decrypt(&ct, &key).unwrap();
236        assert_eq!(pt, plaintext);
237    }
238
239    #[test]
240    fn decrypt_rejects_tampered_ciphertext() {
241        let key = [0x42u8; 32];
242        let plaintext = b"hello noise world";
243        let mut ct = encrypt(plaintext, &key);
244        if let Some(b) = ct.get_mut(0) {
245            *b ^= 0xFF;
246        }
247        assert!(decrypt(&ct, &key).is_none());
248    }
249
250    #[test]
251    fn decrypt_with_wrong_key_fails() {
252        let key = [0x42u8; 32];
253        let wrong_key = [0x43u8; 32];
254        let plaintext = b"secret";
255        let ct = encrypt(plaintext, &key);
256        assert!(decrypt(&ct, &wrong_key).is_none());
257    }
258
259    #[test]
260    fn encrypt_large_message() {
261        let key = [0x42u8; 32];
262        let plaintext = vec![0xAAu8; 10_000];
263        let ct = encrypt(&plaintext, &key);
264        let pt = decrypt(&ct, &key).unwrap();
265        assert_eq!(pt, plaintext);
266    }
267}