ts_noise: factor out Noise implementation

Cargo.toml

@@ -26,6 +26,7 @@ members = [
     "ts_netstack_smoltcp_core",
     "ts_netstack_smoltcp_socket",
     "ts_nodecapability",
+    "ts_noise",
     "ts_overlay_router",
     "ts_packet",
     "ts_packetfilter",
@@ -59,6 +60,7 @@ rust-version = "1.91.0"
 
 [workspace.dependencies]
 aead = { version = "0.5", features = ["std"] }
+aws-lc-rs = "1.17.0"
 base64 = "0.22"
 blake2 = "0.10"
 bytes = "1"
@@ -125,6 +127,7 @@ ts_netstack_smoltcp = { path = "ts_netstack_smoltcp", version = "0.2.0" }
 ts_netstack_smoltcp_core = { path = "ts_netstack_smoltcp_core", version = "0.2.0" }
 ts_netstack_smoltcp_socket = { path = "ts_netstack_smoltcp_socket", version = "0.2.0" }
 ts_nodecapability = { path = "ts_nodecapability", version = "0.2.0" }
+ts_noise = { path = "ts_noise", version = "0.2.0" }
 ts_overlay_router = { path = "ts_overlay_router", version = "0.2.0" }
 ts_packet = { path = "ts_packet", version = "0.2.0" }
 ts_packetfilter = { path = "ts_packetfilter", version = "0.2.0" }

ts_noise/Cargo.toml

@@ -0,0 +1,28 @@
+[package]
+name = "ts_noise"
+edition.workspace = true
+license.workspace = true
+publish.workspace = true
+version.workspace = true
+repository.workspace = true
+rust-version.workspace = true
+
+[dependencies]
+ts_keys.workspace = true
+ts_packet.workspace = true
+ts_time.workspace = true
+
+aws-lc-rs.workspace = true
+aead.workspace = true
+blake2.workspace = true
+chacha20poly1305.workspace = true
+hkdf.workspace = true
+itertools.workspace = true
+rand.workspace = true
+tracing.workspace = true
+x25519-dalek.workspace = true
+zerocopy.workspace = true
+
+
+[lints]
+workspace = true

ts_noise/src/core.rs

@@ -0,0 +1,258 @@
+//! Supporting machinery for executing Noise handshakes.
+//!
+//! This is not a general-purpose Noise protocol library. The provided functionality is
+//! sufficient to execute the two protocols that Tailscale cares about (IK and IKpsk2).
+//!
+//! This module only provides the primitive operations found in Noise handshakes. It is
+//! the caller's responsibility to chain the primitives together correctly to produce the
+//! desired handshake pattern.
+//!
+//! This module uses typestates to make some invalid sequences of operations compile-time
+//! errors, in an effort to make protocol construction errors more obvious. This does not
+//! cover all possible invalid sequences, it is still the caller's responsibility to ensure
+//! that their sequence correctly reflects the desired handshake pattern.
+
+use aws_lc_rs::aead::{
+    Aad, BoundKey, CHACHA20_POLY1305, LessSafeKey, Nonce, OpeningKey, SealingKey, UnboundKey,
+    nonce_sequence::{Counter64, Counter64Builder},
+};
+use blake2::{Blake2s256, Digest};
+use hkdf::SimpleHkdf;
+use zerocopy::IntoBytes;
+
+/// Initialize a ChaCha20Poly1305 cipher with the given key.
+fn must_cipher(key: [u8; 32]) -> LessSafeKey {
+    // Key construction only fails if the key is the wrong length for the algorithm.
+    let k = UnboundKey::new(&CHACHA20_POLY1305, &key).unwrap();
+    LessSafeKey::new(k)
+}
+
+/// Use HKDF to derive two 32-byte values.
+fn must_hkdf2(chaining_key: &[u8; 32], key: &[u8]) -> ([u8; 32], [u8; 32]) {
+    let kdf = SimpleHkdf::<Blake2s256>::new(Some(chaining_key), key);
+    let mut expanded = [0; 64];
+    // Expansion only fails if you request more bytes than the KDF can provide. This KDF can always
+    // provide 64 bytes.
+    kdf.expand(&[], &mut expanded).unwrap();
+    (
+        expanded[..32].try_into().unwrap(),
+        expanded[32..].try_into().unwrap(),
+    )
+}
+
+/// Use HKDF to derive three 32-byte values.
+fn must_hkdf3(chaining_key: &[u8; 32], key: &[u8]) -> ([u8; 32], [u8; 32], [u8; 32]) {
+    let kdf = SimpleHkdf::<Blake2s256>::new(Some(chaining_key), key);
+    let mut expanded = [0; 96];
+    // Expansion only fails if you request more bytes than the KDF can provide. This KDF can always
+    // provide 96 bytes.
+    kdf.expand(&[], &mut expanded).unwrap();
+    (
+        expanded[..32].try_into().unwrap(),
+        expanded[32..64].try_into().unwrap(),
+        expanded[64..].try_into().unwrap(),
+    )
+}
+
+/// A symmetric session.
+///
+/// The sending key is eagerly bound to a nonce sequence, whereas the receiving key is a plain
+/// UnboundKey due to incoming messages needing different nonce handling in different contexts.
+pub struct Session {
+    /// The key to send data.
+    send: SealingKey<Counter64>,
+    /// The key to receive data.
+    recv: UnboundKey,
+    max_recv_nonce: u64,
+}
+
+/// Base Noise handshake state.
+#[derive(Clone)]
+pub struct State {
+    hash: [u8; 32],
+    chaining_key: [u8; 32],
+}
+
+impl State {
+    /// Initialize a new Noise handshake.
+    ///
+    /// `protocol_name` is the Noise protocol name as specified
+    /// in https://noiseprotocol.org/noise.html#protocol-names-and-modifiers.
+    pub fn new(protocol_name: &[u8]) -> State {
+        let init = Blake2s256::digest(protocol_name);
+        State {
+            hash: init.into(),
+            chaining_key: init.into(),
+        }
+    }
+
+    /// Mix data into the handshake state.
+    ///
+    /// This is the MixHash() operation in the Noise spec.
+    #[inline]
+    pub fn mix_hash(mut self, data: &[u8]) -> Self {
+        self.mix_hash_gather(&[data])
+    }
+
+    /// Like mix_hash, but the data can be provided in multiple non-contiguous pieces.
+    pub fn mix_hash_gather(mut self, data: &[&[u8]]) -> Self {
+        let mut h = Blake2s256::new_with_prefix(self.hash);
+        for d in data {
+            h.update(d);
+        }
+        h.finalize_into(self.hash.as_mut_bytes().into());
+        self
+    }
+
+    /// Mix a public key into the handshake state.
+    ///
+    /// This should only be used to mix in the ephemeral public key in psk handshake variants,
+    /// in accordance with sections 9.2 and 9.3 of the Noise spec. Use [`State::mix_dh`] if
+    /// you're looking to MixKey the result of an X25519 operation.
+    ///
+    /// This is the MixKey() operation in the Noise spec.
+    pub fn mix_key(self, public: &x25519_dalek::PublicKey) -> State {
+        let (ck, _) = must_hkdf2(&self.chaining_key, public.as_ref());
+        State {
+            hash: self.hash,
+            chaining_key: ck,
+        }
+    }
+
+    /// Perform an X25519 operation, and mix the result into the handshake state.
+    ///
+    /// This is the MixKey(DH(private, public)) operation in the Noise spec.
+    pub fn mix_dh(
+        self,
+        private: &x25519_dalek::StaticSecret,
+        public: &x25519_dalek::PublicKey,
+    ) -> StateWithAEAD {
+        let shared = private.diffie_hellman(public);
+        let (ck, k) = must_hkdf2(&self.chaining_key, shared.as_ref());
+        StateWithAEAD {
+            state: State {
+                hash: self.hash,
+                chaining_key: ck,
+            },
+            aead: must_cipher(k),
+        }
+    }
+
+    /// Finalize the handshake as the initiator role.
+    ///
+    /// This is the Split() operation in the Noise spec.
+    pub fn finish_as_initiator(self, max_nonce: u64) -> Session {
+        let (init_to_resp, resp_to_init) = must_hkdf2(&self.chaining_key, &[]);
+
+        let recv = UnboundKey::new(&CHACHA20_POLY1305, &resp_to_init).unwrap();
+
+        let send = UnboundKey::new(&CHACHA20_POLY1305, &init_to_resp).unwrap();
+        let send = SealingKey::new(send, Counter64Builder::new().limit(max_nonce).build());
+
+        Session {
+            send,
+            recv,
+            max_recv_nonce: max_nonce,
+        }
+    }
+
+    /// Finalize the handshake as the responder role.
+    ///
+    /// This is the Split() operation in the Noise spec.
+    pub fn finish_as_responder(self, max_nonce: u64) -> Session {
+        let (init_to_resp, resp_to_init) = must_hkdf2(&self.chaining_key, &[]);
+
+        let recv = UnboundKey::new(&CHACHA20_POLY1305, &init_to_resp).unwrap();
+
+        let send = UnboundKey::new(&CHACHA20_POLY1305, &resp_to_init).unwrap();
+        let send = SealingKey::new(send, Counter64Builder::new().limit(max_nonce).build());
+
+        Session {
+            send,
+            recv,
+            max_recv_nonce: max_nonce,
+        }
+    }
+}
+
+/// Noise handshake state when AEAD operations are available.
+///
+/// For the handshake patterns we care about, when the handshake is in this state there are
+/// only two valid ways to continue:
+///
+/// - Perform an AEAD operation ([`StateWithAEAD::seal`] or [`StateWithAEAD::open`]), which
+///   consumes the AEAD and returns a plain [`State`].
+/// - Mix additional key material into the handshake ([`StateWithAEAD::mix_key`] or
+///   [`StateWithAEAD::mix_psk`], which returns an updated [`StateWithAEAD`].
+pub struct StateWithAEAD {
+    state: State,
+    aead: LessSafeKey,
+}
+
+impl StateWithAEAD {
+    /// Perform an X25519 operation, and mix the result into the handshake state.
+    ///
+    /// This is the MixKey(DH(private, public)) operation in the Noise spec.
+    pub fn mix_dh(
+        self,
+        private: &x25519_dalek::StaticSecret,
+        public: &x25519_dalek::PublicKey,
+    ) -> StateWithAEAD {
+        self.state.mix_dh(private, public)
+    }
+
+    /// Mix a pre-shared symmetric key into the handshake state.
+    ///
+    /// This is the MixKeyAndHash() operation in the Noise spec.
+    pub fn mix_psk(self, psk: &[u8; 32]) -> StateWithAEAD {
+        let (ck, h, k) = must_hkdf3(&self.state.chaining_key, psk);
+        StateWithAEAD {
+            state: State {
+                hash: self.state.hash,
+                chaining_key: ck,
+            }
+            .mix_hash(&h),
+            aead: must_cipher(k),
+        }
+    }
+
+    /// Seal `cleartext` in place.
+    ///
+    /// `cleartext` overwritten with ciphertext, and the authentication tag is written to `tag`.
+    ///
+    /// # Panics
+    ///
+    /// If `tag` is not exactly 16 bytes.
+    pub fn seal(self, cleartext: &mut [u8], tag: &mut [u8]) -> State {
+        assert_eq!(tag.len(), 16, "tag must be exactly 16 bytes");
+        let nonce = Nonce::assume_unique_for_key([0; 12]);
+        self.aead
+            .seal_in_place_scatter(nonce, Aad::from(&self.state.hash), cleartext, &[], tag)
+            .unwrap();
+        self.state.mix_hash_gather(&[cleartext, tag])
+    }
+
+    /// Decrypt `ciphertext_and_tag` in place and return the cleartext portion of the slice.
+    ///
+    /// Returns None if decryption fails.
+    ///
+    /// This is the DecryptAndHash() operation in the Noise spec.
+    pub fn open(self, ciphertext: &mut [u8], tag: &[u8]) -> Option<State> {
+        assert_eq!(tag.len(), 16, "tag must be exactly 16 bytes");
+
+        // On successful decryption, we have to mix the ciphertext into the handshake state.
+        // This pairs awkwardly with the in place crypto we're doing, where a successful decryption
+        // overwrites the ciphertext.
+        // Instead, update the handshake hash before the decryption attempt. This is okay because
+        // we discard the handshake state entirely on decryption failure.
+        let hash = self.state.hash.clone();
+        let state = self.state.mix_hash_gather(&[ciphertext, tag]);
+
+        let nonce = Nonce::assume_unique_for_key([0; 12]);
+        self.aead
+            .open_in_place_separate_tag(nonce, Aad::from(&hash), tag, ciphertext)
+            .ok()?;
+
+        Some(state)
+    }
+}