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",
@@ -99,9 +100,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.3.3" }
@@ -125,6 +127,7 @@ ts_netstack_smoltcp = { path = "ts_netstack_smoltcp", version = "0.3.3" }
 ts_netstack_smoltcp_core = { path = "ts_netstack_smoltcp_core", version = "0.3.3" }
 ts_netstack_smoltcp_socket = { path = "ts_netstack_smoltcp_socket", version = "0.3.3" }
 ts_nodecapability = { path = "ts_nodecapability", version = "0.3.3" }
+ts_noise = { path = "ts_noise", version = "0.3.3" }
 ts_overlay_router = { path = "ts_overlay_router", version = "0.3.3" }
 ts_packet = { path = "ts_packet", version = "0.3.3" }
 ts_packetfilter = { path = "ts_packetfilter", version = "0.3.3" }

ts_control/src/control_dialer.rs

@@ -255,7 +255,9 @@ where
         CapabilityVersion::CURRENT,
     );
 
-    let conn = crate::tokio::upgrade_ts2021(url, &init_msg, handshake, h1_client).await?;
+    let conn =
+        crate::tokio::upgrade_ts2021(url, &init_msg, handshake, machine_keys.private, h1_client)
+            .await?;
     let conn = crate::tokio::read_challenge_packet(conn).await?;
 
     let h2_conn = ts_http_util::http2::connect(conn).await?;

ts_control/src/tokio/connect.rs

@@ -5,7 +5,7 @@ use bytes::Bytes;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite};
 use ts_capabilityversion::CapabilityVersion;
 use ts_http_util::{BytesBody, ClientExt, EmptyBody, HeaderName, HeaderValue, Http2, ResponseExt};
-use ts_keys::{MachineKeyPair, MachinePublicKey};
+use ts_keys::{MachineKeyPair, MachinePrivateKey, MachinePublicKey};
 use url::Url;
 use zerocopy::network_endian::U32;
 
@@ -166,7 +166,14 @@ pub async fn connect(
         CapabilityVersion::CURRENT,
     );
 
-    let conn = upgrade_ts2021(control_url, &init_msg, handshake, h1_client).await?;
+    let conn = upgrade_ts2021(
+        control_url,
+        &init_msg,
+        handshake,
+        machine_keys.private,
+        h1_client,
+    )
+    .await?;
 
     // The early payload (challenge packet) is optional. The server may send
     // the magic prefix [FF FF FF 'T' 'S'] followed by a JSON challenge, or it
@@ -225,7 +232,8 @@ pub async fn fetch_control_key(control_url: &Url) -> Result<MachinePublicKey, Co
 pub async fn upgrade_ts2021(
     control_url: &Url,
     init_msg: &str,
-    mut handshake: ts_control_noise::Handshake,
+    handshake: ts_control_noise::Handshake,
+    machine_private_key: MachinePrivateKey,
     h1_client: impl ts_http_util::Client<EmptyBody>,
 ) -> Result<impl AsyncRead + AsyncWrite + Unpin + 'static, ConnectionError> {
     let ts2021_url = control_url.join("/ts2021")?;
@@ -251,7 +259,7 @@ pub async fn upgrade_ts2021(
         ConnectionError::Internal(InternalErrorKind::Http)
     })?;
 
-    let conn = handshake.complete(upgraded).await?;
+    let conn = handshake.complete(upgraded, &machine_private_key).await?;
 
     tracing::debug!("upgraded control connection from HTTP/1.1 to TS2021");
 

ts_control_noise/Cargo.toml

@@ -15,6 +15,7 @@ rust-version.workspace = true
 ts_capabilityversion.workspace = true
 ts_hexdump.workspace = true
 ts_keys.workspace = true
+ts_noise.workspace = true
 
 # Unconditionally required dependencies.
 base64.workspace = true
@@ -30,7 +31,6 @@ tokio = { workspace = true, features = ["io-util"] }
 tokio-util = { workspace = true, features = ["codec"] }
 tracing.workspace = true
 zerocopy.workspace = true
-zeroize = "1.8"
 
 [dev-dependencies]
 rand = "0.10"

ts_control_noise/src/codec.rs

@@ -1,35 +1,42 @@
-use std::io::ErrorKind;
+use std::{io::ErrorKind, marker::PhantomData};
 
 use bytes::{Buf, BufMut, BytesMut};
-use noise_protocol::{Cipher, CipherState};
+use chacha20poly1305::{AeadInPlace, ChaCha20Poly1305, Key, KeyInit, Nonce};
 use tokio_util::codec::{Decoder, Encoder};
-use zerocopy::{IntoBytes, TryCastError, TryFromBytes, U16};
+use ts_noise::core::Session;
+use zerocopy::{IntoBytes, TryCastError, TryFromBytes, network_endian::U16};
 
 use crate::messages::{Header, MessageType};
 
 /// The maximum wire size of a message to control over noise.
 pub const MAX_MESSAGE_SIZE: usize = 4096;
 
+/// Overhead required by the AEAD's authentication tag.
+const AEAD_OVERHEAD: usize = 16;
+
+/// Maximum size of a data chunk, without the per-message overhead.
+const MAX_CHUNK_SIZE: usize = MAX_MESSAGE_SIZE - size_of::<Header>() - AEAD_OVERHEAD;
+
+/// Marker type to indicate that a [`Codec`] instance can only be used for receiving, not sending.
+pub struct Rx {}
+
+/// Marker type to indicate that a [`Codec`] instance can only be used for sending, not receiving.
+pub struct Tx {}
+
 /// Control noise codec that uses a different cipher state for the up and down directions.
 ///
 /// Just a wrapper containing two [`Codec`]s, one of which provides [`Encoder`] and the
 /// other [`Decoder`].
-pub struct BiCodec<Tx, Rx>
-where
-    Tx: Cipher,
-    Rx: Cipher,
-{
+pub struct BiCodec {
     /// The transmit codec, used for encoding messages to control.
     pub tx: Codec<Tx>,
     /// The receive codec, used for decoding messages from control.
     pub rx: Codec<Rx>,
 }
 
-impl<B, Tx, Rx> Encoder<B> for BiCodec<Tx, Rx>
+impl<B> Encoder<B> for BiCodec
 where
     B: AsRef<[u8]>,
-    Tx: Cipher,
-    Rx: Cipher,
 {
     type Error = <Codec<Tx> as Encoder<B>>::Error;
 
@@ -38,11 +45,7 @@ where
     }
 }
 
-impl<Tx, Rx> Decoder for BiCodec<Tx, Rx>
-where
-    Tx: Cipher,
-    Rx: Cipher,
-{
+impl Decoder for BiCodec {
     type Item = <Codec<Rx> as Decoder>::Item;
     type Error = <Codec<Rx> as Decoder>::Error;
 
@@ -51,65 +54,93 @@ where
     }
 }
 
+impl From<Session> for BiCodec {
+    fn from(session: Session) -> Self {
+        Self {
+            tx: Codec::<Tx>::from(session.send),
+            rx: Codec::<Rx>::from(session.send),
+        }
+    }
+}
+
 /// Codec supporting encrypting and decrypting data according to the control noise protocol
 /// using the specified cipher state.
-pub struct Codec<C>
-where
-    C: Cipher,
-{
-    /// The cipher state to use to encode and decode message payloads.
-    pub cipher_state: CipherState<C>,
+///
+/// In accordance with Noise session semantics, a particular Codec instance can only be used for
+/// sending or receiving, never both. The type parameter should be [`Tx`] for sending sessions and
+/// [`Rx`] for receiving sessions.
+pub struct Codec<D> {
+    cipher: ChaCha20Poly1305,
+    next_nonce: u64,
+    _phantom: PhantomData<D>,
 }
 
-impl<C> From<CipherState<C>> for Codec<C>
-where
-    C: Cipher,
-{
-    fn from(value: CipherState<C>) -> Self {
+impl<D> Codec<D> {
+    fn next_nonce(&mut self) -> Nonce {
+        assert_ne!(self.next_nonce, u64::MAX);
+        let mut ret = [0; 12];
+        ret[4..].copy_from_slice(&self.next_nonce.to_be_bytes());
+        self.next_nonce += 1;
+        ret.into()
+    }
+}
+
+impl<D> From<Key> for Codec<D> {
+    fn from(key: Key) -> Self {
+        Codec {
+            cipher: ChaCha20Poly1305::new(&key),
+            next_nonce: 0,
+            _phantom: PhantomData,
+        }
+    }
+}
+
+#[cfg(test)]
+impl<D> From<([u8; 32], u64)> for Codec<D> {
+    fn from((key, nonce): ([u8; 32], u64)) -> Self {
         Codec {
-            cipher_state: value,
+            cipher: ChaCha20Poly1305::new(&key.into()),
+            next_nonce: nonce,
+            _phantom: PhantomData,
         }
     }
 }
 
-impl<B, C> Encoder<B> for Codec<C>
+impl<B> Encoder<B> for Codec<Tx>
 where
-    C: Cipher,
     B: AsRef<[u8]>,
 {
     type Error = std::io::Error;
 
     fn encode(&mut self, b: B, dst: &mut BytesMut) -> Result<(), Self::Error> {
         let b = b.as_ref();
-        let max_data_chunk = MAX_MESSAGE_SIZE - (size_of::<Header>() + C::tag_len());
 
-        for chunk in b.chunks(max_data_chunk) {
+        for chunk in b.chunks(MAX_CHUNK_SIZE) {
             let hdr = Header {
                 typ: MessageType::Record,
-                len: U16::new(chunk.len() as u16 + C::tag_len() as u16),
+                len: U16::new(chunk.len() as u16 + AEAD_OVERHEAD as u16),
             };
 
             dst.put(hdr.as_bytes());
 
             let data_start = dst.len();
-
             dst.put(chunk);
-            dst.put_bytes(0, C::tag_len());
 
-            self.cipher_state
-                .encrypt_in_place(&mut dst[data_start..], chunk.len());
+            let nonce = self.next_nonce();
+            let tag = self
+                .cipher
+                .encrypt_in_place_detached(&nonce, &[], &mut dst[data_start..])
+                .unwrap();
+            dst.put(tag.as_ref());
         }
 
         Ok(())
     }
 }
 
-impl<C> Decoder for Codec<C>
-where
-    C: Cipher,
-{
-    type Error = std::io::Error;
+impl Decoder for Codec<Rx> {
     type Item = BytesMut;
+    type Error = std::io::Error;
 
     fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
         let (header, rest_len) = match Header::try_ref_from_prefix(src) {
@@ -131,12 +162,15 @@ where
 
         match header.typ {
             MessageType::Record => {
-                match self.cipher_state.decrypt_in_place(&mut body, len) {
-                    Ok(n) => body.truncate(n),
-                    Err(()) => {
-                        tracing::error!("decryption failed");
-                        return Err(ErrorKind::InvalidData.into());
-                    }
+                let nonce = self.next_nonce();
+                let tag = body.split_off(body.len() - AEAD_OVERHEAD);
+                if self
+                    .cipher
+                    .decrypt_in_place_detached(&nonce, &[], body.as_mut(), tag.as_ref().into())
+                    .is_err()
+                {
+                    tracing::error!("decryption failed");
+                    return Err(ErrorKind::InvalidData.into());
                 }
 
                 Ok(Some(body))
@@ -173,21 +207,11 @@ mod test {
 
     type Cipher = crate::ChaCha20Poly1305BigEndian;
 
-    fn init_codec_pair(key: [u8; 32], nonce: u64) -> (Codec<Cipher>, Codec<Cipher>) {
-        let encrypt_state = CipherState::<Cipher>::new(&key, nonce);
-        let decrypt_state = encrypt_state.clone();
-
-        (
-            Codec {
-                cipher_state: encrypt_state,
-            },
-            Codec {
-                cipher_state: decrypt_state,
-            },
-        )
+    fn init_codec_pair(key: [u8; 32], nonce: u64) -> (Codec<Tx>, Codec<Rx>) {
+        ((key, nonce).into(), (key, nonce).into())
     }
 
-    fn rand_codec_pair() -> (Codec<Cipher>, Codec<Cipher>) {
+    fn rand_codec_pair() -> (Codec<Tx>, Codec<Rx>) {
         init_codec_pair(rand::random(), rand::random())
     }