1use chacha20poly1305::{
12 aead::{Aead, KeyInit},
13 ChaCha20Poly1305, Nonce,
14};
15use sha2::{Digest, Sha256};
16use snow::{params::NoiseParams, Builder, HandshakeState, TransportState};
17use std::sync::OnceLock;
18use x25519_dalek::{PublicKey, StaticSecret};
19
20pub const NOISE_PATTERN: &str = "Noise_NK_25519_ChaChaPoly_SHA256";
22
23fn noise_params() -> &'static NoiseParams {
24 static PARAMS: OnceLock<NoiseParams> = OnceLock::new();
25 PARAMS.get_or_init(|| NOISE_PATTERN.parse().expect("Noise_NK params"))
26}
27
28#[derive(Clone, Copy, Debug, PartialEq, Eq)]
30pub struct PublicKeyBytes(pub [u8; 32]);
31
32impl PublicKeyBytes {
33 pub fn as_bytes(&self) -> &[u8; 32] {
34 &self.0
35 }
36}
37
38#[derive(Clone)]
40pub struct RestreamerKeys {
41 c2_secret: StaticSecret,
42 c2_public: PublicKey,
43 hq_recipient_public: PublicKey,
44 hq_recipient_public_prev: Option<PublicKey>,
45}
46
47impl RestreamerKeys {
48 pub fn generate(hq_recipient_public: PublicKeyBytes) -> Self {
49 let mut seed = [0u8; 32];
50 getrandom::getrandom(&mut seed).expect("os rng");
51 Self::from_seeds(seed, hq_recipient_public.0)
52 }
53
54 pub fn from_seeds(c2_seed: [u8; 32], hq_public: [u8; 32]) -> Self {
55 let c2_secret = StaticSecret::from(c2_seed);
56 let c2_public = PublicKey::from(&c2_secret);
57 Self {
58 c2_secret,
59 c2_public,
60 hq_recipient_public: PublicKey::from(hq_public),
61 hq_recipient_public_prev: None,
62 }
63 }
64
65 pub fn upstream_public_keys(&self) -> (PublicKeyBytes, PublicKeyBytes) {
66 (
67 PublicKeyBytes(*self.c2_public.as_bytes()),
68 PublicKeyBytes(*self.hq_recipient_public.as_bytes()),
69 )
70 }
71
72 pub fn rotate_hq_recipient(&mut self, new_public: PublicKeyBytes) {
73 self.hq_recipient_public_prev = Some(self.hq_recipient_public);
74 self.hq_recipient_public = PublicKey::from(new_public.0);
75 }
76
77 pub fn end_rotation_window(&mut self) {
78 self.hq_recipient_public_prev = None;
79 }
80
81 pub fn hq_recipient_public(&self) -> PublicKeyBytes {
82 PublicKeyBytes(*self.hq_recipient_public.as_bytes())
83 }
84
85 pub fn previous_hq_recipient_public(&self) -> Option<PublicKeyBytes> {
86 self.hq_recipient_public_prev
87 .map(|k| PublicKeyBytes(*k.as_bytes()))
88 }
89
90 pub fn c2_public(&self) -> PublicKeyBytes {
91 PublicKeyBytes(*self.c2_public.as_bytes())
92 }
93}
94
95#[derive(Clone)]
97pub struct HqRecipientSecret {
98 private: [u8; 32],
99 public: [u8; 32],
100}
101
102impl HqRecipientSecret {
103 pub fn generate() -> Self {
104 let mut seed = [0u8; 32];
105 getrandom::getrandom(&mut seed).expect("os rng");
106 Self::from_seed(seed)
107 }
108
109 pub fn from_seed(seed: [u8; 32]) -> Self {
110 let secret = StaticSecret::from(seed);
111 let public = PublicKey::from(&secret);
112 Self {
113 private: seed,
114 public: *public.as_bytes(),
115 }
116 }
117
118 pub fn public_bytes(&self) -> PublicKeyBytes {
119 PublicKeyBytes(self.public)
120 }
121
122 pub fn seed_bytes(&self) -> [u8; 32] {
124 self.private
125 }
126}
127
128pub struct NoiseInitiator {
130 state: HandshakeState,
131}
132
133impl NoiseInitiator {
134 pub fn new(hq_recipient_public: PublicKeyBytes) -> Result<Self, CryptoError> {
135 let state = Builder::new(noise_params().clone())
136 .remote_public_key(&hq_recipient_public.0)
137 .map_err(noise_err)?
138 .build_initiator()
139 .map_err(noise_err)?;
140 Ok(Self { state })
141 }
142
143 pub fn write_message1(&mut self) -> Result<Vec<u8>, CryptoError> {
145 let mut buf = vec![0u8; 64];
146 let n = self.state.write_message(&[], &mut buf).map_err(noise_err)?;
147 buf.truncate(n);
148 Ok(buf)
149 }
150
151 pub fn read_message2(self, msg: &[u8]) -> Result<NoiseTransport, CryptoError> {
153 let mut state = self.state;
154 let mut payload = [0u8; 64];
155 let _ = state.read_message(msg, &mut payload).map_err(noise_err)?;
156 let transport = state.into_transport_mode().map_err(noise_err)?;
157 Ok(NoiseTransport { state: transport })
158 }
159}
160
161pub struct NoiseResponder {
163 state: HandshakeState,
164}
165
166impl NoiseResponder {
167 pub fn new(hq: &HqRecipientSecret) -> Result<Self, CryptoError> {
168 let state = Builder::new(noise_params().clone())
169 .local_private_key(&hq.private)
170 .map_err(noise_err)?
171 .build_responder()
172 .map_err(noise_err)?;
173 Ok(Self { state })
174 }
175
176 pub fn complete(self, message1: &[u8]) -> Result<(Vec<u8>, NoiseTransport), CryptoError> {
178 let mut state = self.state;
179 let mut payload = [0u8; 64];
180 let _ = state.read_message(message1, &mut payload).map_err(noise_err)?;
181 let mut msg2 = vec![0u8; 64];
182 let n = state.write_message(&[], &mut msg2).map_err(noise_err)?;
183 msg2.truncate(n);
184 let transport = state.into_transport_mode().map_err(noise_err)?;
185 Ok((msg2, NoiseTransport { state: transport }))
186 }
187}
188
189pub struct NoiseTransport {
191 state: TransportState,
192}
193
194impl NoiseTransport {
195 pub fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>, CryptoError> {
196 let mut buf = vec![0u8; plaintext.len() + 16];
197 let n = self.state.write_message(plaintext, &mut buf).map_err(noise_err)?;
198 buf.truncate(n);
199 Ok(buf)
200 }
201
202 pub fn decrypt(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>, CryptoError> {
203 let mut buf = vec![0u8; ciphertext.len()];
204 let n = self.state.read_message(ciphertext, &mut buf).map_err(noise_err)?;
205 buf.truncate(n);
206 Ok(buf)
207 }
208}
209
210pub struct NestedSeal {
212 pub outer: Vec<u8>,
214 pub inner_ciphertext: Vec<u8>,
216 pub hq_transport: NoiseTransport,
218 pub node_transport: NoiseTransport,
220 pub handshake_msg1: Vec<u8>,
221 pub handshake_msg2: Vec<u8>,
222}
223
224pub fn noise_session_seal(
226 keys: &RestreamerKeys,
227 hq: &HqRecipientSecret,
228 plaintext: &[u8],
229) -> Result<NestedSeal, CryptoError> {
230 let mut initiator = NoiseInitiator::new(keys.hq_recipient_public())?;
231 let msg1 = initiator.write_message1()?;
232 let responder = NoiseResponder::new(hq)?;
233 let (msg2, hq_transport) = responder.complete(&msg1)?;
234 let mut node_transport = initiator.read_message2(&msg2)?;
235 let inner_ciphertext = node_transport.encrypt(plaintext)?;
236 let outer = outer_seal(&keys.c2_public, &inner_ciphertext)?;
237 Ok(NestedSeal {
238 outer,
239 inner_ciphertext,
240 hq_transport,
241 node_transport,
242 handshake_msg1: msg1,
243 handshake_msg2: msg2,
244 })
245}
246
247pub fn core_open_inner(
249 transport: &mut NoiseTransport,
250 ciphertext: &[u8],
251) -> Result<Vec<u8>, CryptoError> {
252 transport.decrypt(ciphertext)
253}
254
255pub fn node_seal_outer(c2_public: PublicKeyBytes, inner: &[u8]) -> Result<Vec<u8>, CryptoError> {
257 outer_seal(&PublicKey::from(c2_public.0), inner)
258}
259
260pub fn c2_strip_outer(keys: &RestreamerKeys, sealed: &[u8]) -> Result<Vec<u8>, CryptoError> {
262 outer_open(&keys.c2_secret, sealed)
263}
264
265pub fn c2_cannot_complete_as_hq(keys: &RestreamerKeys, handshake_msg1: &[u8]) -> bool {
267 let fake = HqRecipientSecret {
268 private: keys.c2_secret.to_bytes(),
269 public: *PublicKey::from(&keys.c2_secret).as_bytes(),
270 };
271 NoiseResponder::new(&fake)
272 .and_then(|r| r.complete(handshake_msg1))
273 .is_err()
274}
275
276pub fn c2_cannot_open_inner(keys: &RestreamerKeys, inner_ciphertext: &[u8]) -> bool {
279 let _ = keys;
283 outer_open(&keys.c2_secret, inner_ciphertext).is_err()
285}
286
287fn outer_seal(recipient: &PublicKey, plaintext: &[u8]) -> Result<Vec<u8>, CryptoError> {
288 let mut eph_seed = [0u8; 32];
289 getrandom::getrandom(&mut eph_seed).map_err(|e| CryptoError::Rng(e.to_string()))?;
290 let eph = StaticSecret::from(eph_seed);
291 let eph_pub = PublicKey::from(&eph);
292 let shared = eph.diffie_hellman(recipient);
293 let key = kdf(shared.as_bytes());
294 let cipher = ChaCha20Poly1305::new_from_slice(&key).map_err(|_| CryptoError::BadKey)?;
295 let mut nonce_bytes = [0u8; 12];
296 getrandom::getrandom(&mut nonce_bytes).map_err(|e| CryptoError::Rng(e.to_string()))?;
297 let nonce = Nonce::from_slice(&nonce_bytes);
298 let ct = cipher
299 .encrypt(nonce, plaintext)
300 .map_err(|_| CryptoError::Encrypt)?;
301 let mut out = Vec::with_capacity(32 + 12 + ct.len());
302 out.extend_from_slice(eph_pub.as_bytes());
303 out.extend_from_slice(&nonce_bytes);
304 out.extend_from_slice(&ct);
305 Ok(out)
306}
307
308fn outer_open(secret: &StaticSecret, sealed: &[u8]) -> Result<Vec<u8>, CryptoError> {
309 if sealed.len() < 32 + 12 + 16 {
310 return Err(CryptoError::Truncated);
311 }
312 let mut eph_bytes = [0u8; 32];
313 eph_bytes.copy_from_slice(&sealed[..32]);
314 let eph_pub = PublicKey::from(eph_bytes);
315 let nonce = Nonce::from_slice(&sealed[32..44]);
316 let ct = &sealed[44..];
317 let shared = secret.diffie_hellman(&eph_pub);
318 let key = kdf(shared.as_bytes());
319 let cipher = ChaCha20Poly1305::new_from_slice(&key).map_err(|_| CryptoError::BadKey)?;
320 cipher
321 .decrypt(nonce, ct)
322 .map_err(|_| CryptoError::Decrypt)
323}
324
325fn kdf(shared: &[u8]) -> [u8; 32] {
326 let mut hasher = Sha256::new();
327 hasher.update(b"hatchery-restreamer-outer-v1");
328 hasher.update(shared);
329 let digest = hasher.finalize();
330 let mut out = [0u8; 32];
331 out.copy_from_slice(&digest);
332 out
333}
334
335fn noise_err(err: snow::Error) -> CryptoError {
336 CryptoError::Noise(err.to_string())
337}
338
339#[derive(Clone, Debug, PartialEq, Eq)]
340pub enum CryptoError {
341 Rng(String),
342 BadKey,
343 Encrypt,
344 Decrypt,
345 Truncated,
346 Noise(String),
347}
348
349impl std::fmt::Display for CryptoError {
350 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
351 match self {
352 Self::Rng(msg) => write!(f, "rng: {msg}"),
353 Self::BadKey => write!(f, "bad key"),
354 Self::Encrypt => write!(f, "encrypt failed"),
355 Self::Decrypt => write!(f, "decrypt failed"),
356 Self::Truncated => write!(f, "truncated sealed box"),
357 Self::Noise(msg) => write!(f, "noise: {msg}"),
358 }
359 }
360}
361
362impl std::error::Error for CryptoError {}
363
364#[cfg(test)]
365mod tests {
366 use super::*;
367
368 #[test]
369 fn noise_nk_handshake_and_transport() {
370 let hq = HqRecipientSecret::from_seed([7u8; 32]);
371 let mut initiator = NoiseInitiator::new(hq.public_bytes()).unwrap();
372 let msg1 = initiator.write_message1().unwrap();
373 let responder = NoiseResponder::new(&hq).unwrap();
374 let (msg2, mut hq_tx) = responder.complete(&msg1).unwrap();
375 let mut node_tx = initiator.read_message2(&msg2).unwrap();
376 let ct = node_tx.encrypt(b"agent-message-body").unwrap();
377 assert!(!ct.windows(b"agent-message-body".len()).any(|w| w == b"agent-message-body"));
378 assert_eq!(hq_tx.decrypt(&ct).unwrap(), b"agent-message-body");
379 }
380
381 #[test]
382 fn nested_roundtrip_and_c2_blind_to_inner() {
383 let hq = HqRecipientSecret::from_seed([7u8; 32]);
384 let keys = RestreamerKeys::from_seeds([3u8; 32], hq.public_bytes().0);
385 let (c2_pub, hq_pub) = keys.upstream_public_keys();
386 assert_eq!(c2_pub, keys.c2_public());
387 assert_eq!(hq_pub, hq.public_bytes());
388
389 let plain = b"agent-message-body";
390 let sealed = noise_session_seal(&keys, &hq, plain).unwrap();
391 let inner = c2_strip_outer(&keys, &sealed.outer).unwrap();
392 assert_eq!(inner, sealed.inner_ciphertext);
393 assert!(c2_cannot_open_inner(&keys, &inner));
394 assert!(c2_cannot_complete_as_hq(&keys, &sealed.handshake_msg1));
395 assert!(!inner.windows(plain.len()).any(|w| w == plain));
396 let mut hq_tx = sealed.hq_transport;
397 assert_eq!(core_open_inner(&mut hq_tx, &inner).unwrap(), plain);
398 }
399
400 #[test]
401 fn tampered_inner_rejected_by_core() {
402 let hq = HqRecipientSecret::from_seed([9u8; 32]);
403 let keys = RestreamerKeys::from_seeds([1u8; 32], hq.public_bytes().0);
404 let sealed = noise_session_seal(&keys, &hq, b"secret").unwrap();
405 let mut inner = c2_strip_outer(&keys, &sealed.outer).unwrap();
406 let last = inner.len() - 1;
407 inner[last] ^= 0x5a;
408 let mut hq_tx = sealed.hq_transport;
409 assert!(core_open_inner(&mut hq_tx, &inner).is_err());
410 }
411
412 #[test]
413 fn rotation_window_accepts_new_key() {
414 let hq_old = HqRecipientSecret::from_seed([2u8; 32]);
415 let hq_new = HqRecipientSecret::from_seed([4u8; 32]);
416 let mut keys = RestreamerKeys::from_seeds([5u8; 32], hq_old.public_bytes().0);
417 keys.rotate_hq_recipient(hq_new.public_bytes());
418 let sealed = noise_session_seal(&keys, &hq_new, b"after-rotate").unwrap();
419 let inner = c2_strip_outer(&keys, &sealed.outer).unwrap();
420 let mut hq_tx = sealed.hq_transport;
421 assert_eq!(core_open_inner(&mut hq_tx, &inner).unwrap(), b"after-rotate");
422 assert!(NoiseResponder::new(&hq_old)
424 .unwrap()
425 .complete(&sealed.handshake_msg1)
426 .is_err());
427 keys.end_rotation_window();
428 assert!(keys.previous_hq_recipient_public().is_none());
429 }
430
431 #[test]
432 fn wrong_recipient_cannot_complete_handshake() {
433 let hq = HqRecipientSecret::from_seed([1u8; 32]);
434 let other = HqRecipientSecret::from_seed([2u8; 32]);
435 let mut initiator = NoiseInitiator::new(hq.public_bytes()).unwrap();
436 let msg1 = initiator.write_message1().unwrap();
437 let responder = NoiseResponder::new(&other).unwrap();
438 assert!(responder.complete(&msg1).is_err());
439 }
440}