ts_noise: factor out Noise implementations

Also switch to use aws-lc-rs. We already have to pull in aws-lc-rs
for rustls, so avoid using rustcrypto as well for binary size.

Signed-off-by: David Anderson <danderson@tailscale.com>
Change-Id: I690c8c138843750e0486ff9b1cf36df26a6a6964

Author: danderson

Cargo.lock

@@ -26,6 +26,7 @@ members = [
     "ts_netstack_smoltcp_core",
     "ts_netstack_smoltcp_socket",
     "ts_nodecapability",
+    "ts_noise",
     "ts_overlay_router",
     "ts_packet",
     "ts_packetfilter",
@@ -100,9 +101,10 @@ tokio-util = { version = "0.7", default-features = false }
 tracing = { version = "0.1", default-features = false, features = ["attributes"] }
 tracing-test = "0.2"
 url = { version = "2", default-features = false }
-x25519-dalek = { version = "3.0.0-pre.6", features = ["getrandom", "reusable_secrets", "static_secrets"] }
+x25519-dalek = { version = "3.0.0-pre.6", features = ["getrandom", "reusable_secrets", "static_secrets", "zeroize"] }
 yoke = { version = "0.8", default-features = false }
 zerocopy = { version = "0.8", features = ["derive"] }
+zeroize = { version = "1.8.2", features = ["zeroize_derive"] }
 
 # local workspace deps
 tailscale = { path = ".", version = "0.2.0" }
@@ -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" }

Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "ts_noise"
+edition.workspace = true
+license.workspace = true
+publish.workspace = true
+version.workspace = true
+repository.workspace = true
+rust-version.workspace = true
+
+[dependencies]
+aead.workspace = true
+blake2.workspace = true
+chacha20poly1305.workspace = true
+hkdf.workspace = true
+itertools.workspace = true
+x25519-dalek.workspace = true
+zerocopy.workspace = true
+zeroize.workspace = true
+
+[dev-dependencies]
+hex.workspace = true
+
+[lints]
+workspace = true

ts_noise/Cargo.toml

@@ -0,0 +1,243 @@
+//! 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 aead::AeadInPlace;
+use blake2::{Blake2s256, Digest};
+use chacha20poly1305::{ChaCha20Poly1305, KeyInit};
+use hkdf::SimpleHkdf;
+use zerocopy::IntoBytes;
+use zeroize::ZeroizeOnDrop;
+
+/// Initialize a ChaCha20Poly1305 cipher with the given key.
+fn must_cipher(key: [u8; 32]) -> ChaCha20Poly1305 {
+    ChaCha20Poly1305::new(&key.into())
+}
+
+/// 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.
+pub struct Session {
+    /// The key to send data.
+    pub send: chacha20poly1305::Key,
+    /// The key to receive data.
+    pub recv: chacha20poly1305::Key,
+}
+
+/// Base Noise handshake state.
+#[derive(Clone, ZeroizeOnDrop)]
+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(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 MixHash(ephemeral)+MixKey(ephemeral) operations in the Noise spec.
+    pub fn mix_hash_and_key(mut self, public: &x25519_dalek::PublicKey) -> State {
+        self = self.mix_hash(public.as_ref());
+        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) -> Session {
+        let (send, recv) = must_hkdf2(&self.chaining_key, &[]);
+
+        Session {
+            send: send.into(),
+            recv: recv.into(),
+        }
+    }
+
+    /// Finalize the handshake as the responder role.
+    ///
+    /// This is the Split() operation in the Noise spec.
+    pub fn finish_as_responder(self) -> Session {
+        let (recv, send) = must_hkdf2(&self.chaining_key, &[]);
+
+        Session {
+            send: send.into(),
+            recv: recv.into(),
+        }
+    }
+}
+
+/// A pre-shared symmetric key.
+pub type Psk = [u8; 32];
+
+/// The authentication tag of an AEAD ciphertext.
+pub type AeadTag = [u8; 16];
+
+/// Noise handshake state when AEAD operations are available.
+///
+/// For the supported handshake patterns, 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_dh`] or
+///   [`StateWithAead::mix_psk`], which returns an updated [`StateWithAead`].
+pub struct StateWithAead {
+    state: State,
+    aead: ChaCha20Poly1305,
+}
+
+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: &Psk) -> 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`.
+    pub fn seal(self, cleartext: &mut [u8], tag: &mut AeadTag) -> State {
+        let nonce = [0; 12];
+        let res = self
+            .aead
+            .encrypt_in_place_detached(&nonce.into(), &self.state.hash, cleartext)
+            .unwrap();
+        res.write_to(tag).unwrap();
+
+        self.state.mix_hash_gather(&[cleartext, tag])
+    }
+
+    /// Decrypt `ciphertext` in place, authenticating with `tag`.
+    ///
+    /// If successful, `ciphertext` is overwritten with cleartext. Returns `None` if decryption
+    /// fails, implicitly terminating the handshake.
+    ///
+    /// This is the DecryptAndHash() operation in the Noise spec.
+    pub fn open(self, ciphertext: &mut [u8], tag: &AeadTag) -> Option<State> {
+        // 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;
+        let state = self.state.mix_hash_gather(&[ciphertext, tag]);
+
+        let nonce = [0; 12];
+        self.aead
+            .decrypt_in_place_detached(&nonce.into(), &hash, ciphertext, tag.into())
+            .ok()?;
+
+        Some(state)
+    }
+}