ts_keys: rework to use generic types with marker traits

src/lib.rs

@@ -371,9 +371,7 @@ pub mod netstack {
 pub mod keys {
     #[doc(inline)]
     pub use ts_keys::{
-        DiscoKeyPair, DiscoPrivateKey, DiscoPublicKey, MachineKeyPair, MachinePrivateKey,
-        MachinePublicKey, NetworkLockKeyPair, NetworkLockPrivateKey, NetworkLockPublicKey,
-        NodeKeyPair, NodePrivateKey, NodePublicKey, NodeState, PersistState,
+        Disco, KeyPair, Machine, NetworkLock, Node, NodeState, PersistState, PrivateKey, PublicKey,
     };
 }
 

ts_control/src/client/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::{KeyPair, Machine, PublicKey};
 use url::Url;
 use zerocopy::network_endian::U32;
 
@@ -143,8 +143,8 @@ impl From<InternalErrorKind> for crate::InternalErrorKind {
 #[derive(serde::Deserialize, serde::Serialize)]
 #[serde(rename_all = "camelCase")]
 struct ControlPublicKeys {
-    legacy_public_key: MachinePublicKey,
-    public_key: MachinePublicKey,
+    legacy_public_key: PublicKey<Machine>,
+    public_key: PublicKey<Machine>,
 }
 
 impl fmt::Display for ControlPublicKeys {
@@ -160,7 +160,7 @@ impl fmt::Display for ControlPublicKeys {
 #[tracing::instrument(skip_all, fields(%control_url), err)]
 pub async fn connect(
     control_url: &Url,
-    machine_keys: &MachineKeyPair,
+    machine_keys: &KeyPair<Machine>,
 ) -> Result<Http2<BytesBody>, ConnectionError> {
     let h1_client = connect_h1(control_url).await?;
 
@@ -195,15 +195,15 @@ async fn connect_h1(url: &Url) -> Result<ts_http_util::Http1<EmptyBody>, Connect
     }
 }
 
-/// Fetch the control server's [`MachinePublicKey`], which is used to encrypt the Noise connection
-/// to the control server.
+/// Fetch the control server's [machine public key][PublicKey<Machine>], which is
+/// used to encrypt the Noise connection to the control server.
 ///
 /// If the `insecure-keyfetch` feature flag is not enabled, this forces the key url scheme to HTTPS
 /// and validates the server's certificate. This is a critical, load-bearing requirement for
 /// security: if the control server key is MITMed, it's game over. `insecure-keyfetch` is provided
 /// ONLY for integration testing, where it's not practical to issue valid certs.
 #[tracing::instrument(skip_all, fields(%control_url), ret, err, level = "trace")]
-pub async fn fetch_control_key(control_url: &Url) -> Result<MachinePublicKey, ConnectionError> {
+pub async fn fetch_control_key(control_url: &Url) -> Result<PublicKey<Machine>, ConnectionError> {
     let mut key_url = control_url.join("/key")?;
 
     #[cfg(not(feature = "insecure-keyfetch"))]
@@ -242,7 +242,7 @@ pub async fn upgrade_ts2021(
     control_url: &Url,
     init_msg: &str,
     handshake: ts_control_noise::Handshake,
-    machine_key: &MachineKeyPair,
+    machine_key: &KeyPair<Machine>,
     h1_client: impl ts_http_util::Client<EmptyBody>,
 ) -> Result<impl AsyncRead + AsyncWrite + Unpin + 'static, ConnectionError> {
     let ts2021_url = control_url.join("/ts2021")?;

ts_control/src/control_dialer.rs

@@ -9,6 +9,7 @@ use tokio_util::future::FutureExt;
 use ts_bitset::BitsetDyn;
 use ts_capabilityversion::CapabilityVersion;
 use ts_http_util::{BytesBody, Http2};
+use ts_keys::{KeyPair, Machine};
 use url::Url;
 
 use crate::{DialCandidate, DialMode, DialPlan, Error, InternalErrorKind, Operation};
@@ -188,7 +189,7 @@ impl ControlDialer {
     pub async fn full_connect_next(
         &mut self,
         url: &Url,
-        machine_keys: &ts_keys::MachineKeyPair,
+        machine_keys: &KeyPair<Machine>,
     ) -> Result<Http2<BytesBody>, Error> {
         let next = self.next_dialer();
         tracing::trace!(selected_control_dialer = ?next);
@@ -221,7 +222,7 @@ impl ControlDialer {
 /// inner http2 connection.
 pub async fn complete_connection<Io>(
     url: &Url,
-    machine_keys: &ts_keys::MachineKeyPair,
+    machine_keys: &KeyPair<Machine>,
     stream: Io,
 ) -> Result<Http2<BytesBody>, Error>
 where

ts_control/src/node.rs

@@ -6,7 +6,7 @@ use core::{
 
 use chrono::{DateTime, Utc};
 use ts_capabilityversion::CapabilityVersion;
-use ts_keys::{DiscoPublicKey, MachinePublicKey, NodePublicKey};
+use ts_keys::{Disco, Machine, Node as NodeKey, PublicKey};
 
 const LAST_SEEN_FORMAT: &str = "%F %T %Z";
 
@@ -170,15 +170,15 @@ pub struct Node {
     /// The address of the node in the tailnet.
     pub tailnet_address: TailnetAddress,
 
-    /// The node's [`NodePublicKey`].
-    pub node_key: NodePublicKey,
+    /// The node's node public key.
+    pub node_key: PublicKey<NodeKey>,
     /// The node key's expiration.
     pub node_key_expiry: Option<DateTime<Utc>>,
 
-    /// The node's [`MachinePublicKey`], if known.
-    pub machine_key: Option<MachinePublicKey>,
-    /// The node's [`DiscoPublicKey`], if known.
-    pub disco_key: Option<DiscoPublicKey>,
+    /// The node's machine public key, if known.
+    pub machine_key: Option<PublicKey<Machine>>,
+    /// The node's disco public key, if known.
+    pub disco_key: Option<PublicKey<Disco>>,
     /// The signature of the node's public key with the Tailnet Lock signing key, if Tailnet Lock
     /// is enabled and the signature is known.
     pub tailnet_lock_key_signature: Option<Vec<u8>>,
@@ -400,13 +400,13 @@ pub struct NodeUpdate {
     /// The node's capabilities (node caps, not peer caps). If `None`, has not changed.
     pub cap_map: Option<BTreeMap<String, Vec<String>>>,
 
-    /// The node's [`NodePublicKey`]. If `None`, has not changed.
-    pub node_key: Option<NodePublicKey>,
+    /// The node's node public key. If `None`, has not changed.
+    pub node_key: Option<PublicKey<NodeKey>>,
     /// The node key's expiration. If `None`, has not changed.
     pub node_key_expiry: Option<DateTime<Utc>>,
 
-    /// The node's [`DiscoPublicKey`]. If `None`, has not changed.
-    pub disco_key: Option<DiscoPublicKey>,
+    /// The node's disco public key. If `None`, has not changed.
+    pub disco_key: Option<PublicKey<Disco>>,
 
     /// The node's key signature for Tailnet Lock. If `None`, has not changed.
     pub tailnet_lock_key_signature: Option<Vec<u8>>,

ts_control_noise/src/handshake.rs

@@ -3,7 +3,7 @@ use bytes::BytesMut;
 use tokio::io::{AsyncRead, AsyncReadExt};
 use tokio_util::codec::Framed;
 use ts_hexdump::{AsHexExt, Case};
-use ts_keys::{MachineKeyPair, MachinePublicKey};
+use ts_keys::{KeyPair, Machine, PublicKey};
 use ts_noise::ik::SentHandshake;
 use zerocopy::{IntoBytes, TryFromBytes};
 
@@ -19,7 +19,7 @@ type WrappedIo<T> = FramedIo<NoiseFramed<T>, BytesMut>;
 
 /// Noise handshake state.
 pub struct Handshake {
-    state: SentHandshake,
+    state: SentHandshake<Machine>,
 }
 
 impl Handshake {
@@ -31,20 +31,20 @@ impl Handshake {
     /// to the control server in order to start the handshake.
     pub fn initialize(
         prologue: &str,
-        node_machine_key: &MachineKeyPair,
-        control_public_key: &MachinePublicKey,
+        node_machine_key: &KeyPair<Machine>,
+        control_public_key: &PublicKey<Machine>,
         capability_version: ts_capabilityversion::CapabilityVersion,
     ) -> (Self, String) {
-        let mut ciphertext = [0; SentHandshake::INIT_SIZE];
+        let mut ciphertext = [0; SentHandshake::<Machine>::INIT_SIZE];
         let state = SentHandshake::new(
-            node_machine_key.into(),
-            control_public_key.into(),
+            node_machine_key,
+            control_public_key,
             prologue.as_bytes(),
             &mut ciphertext,
         );
         let init_msg = Initiation::new(
             capability_version.into(),
-            SentHandshake::INIT_SIZE as u16,
+            SentHandshake::<Machine>::INIT_SIZE as u16,
             ciphertext,
         );
 
@@ -55,7 +55,7 @@ impl Handshake {
     pub async fn complete<T: AsyncRead + Unpin>(
         mut self,
         mut conn: T,
-        node_machine_key: &MachineKeyPair,
+        node_machine_key: &KeyPair<Machine>,
     ) -> Result<WrappedIo<T>, Error> {
         let mut hdr_bytes = [0u8; 3];
         conn.read_exact(&mut hdr_bytes[..]).await?;
@@ -78,7 +78,10 @@ impl Handshake {
             return Err(Error::BadFormat);
         }
 
-        let session = match self.state.try_finish(&mut packet, node_machine_key.into()) {
+        let session = match self
+            .state
+            .try_finish(&mut packet, &node_machine_key.private)
+        {
             Ok(session) => session,
             Err(state) => {
                 self.state = state;

ts_control_serde/src/netmap.rs

@@ -4,7 +4,7 @@ use core::net::SocketAddr;
 use chrono::{DateTime, Utc};
 use serde::Deserialize;
 use ts_capabilityversion::CapabilityVersion;
-use ts_keys::{DiscoPublicKey, NodePublicKey};
+use ts_keys::{Disco, PublicKey};
 
 use crate::{
     DerpRegionId, DnsConfig, MarshaledSignature,
@@ -50,10 +50,10 @@ pub struct MapRequest<'a> {
     pub keep_alive: bool,
 
     /// The public key of this Tailscale node.
-    pub node_key: NodePublicKey,
+    pub node_key: PublicKey<ts_keys::Node>,
     /// The public key this Tailscale node will use with the Disco protocol to establish direct
     /// connections with peer nodes in the Tailnet.
-    pub disco_key: DiscoPublicKey,
+    pub disco_key: PublicKey<Disco>,
 
     /// If populated, the public key of the node's hardware-backed identity attestation key.
     pub hardware_attestation_key: Option<Vec<u8>>,
@@ -487,14 +487,14 @@ pub struct PeerChange<'a> {
     pub endpoints: Option<Vec<SocketAddr>>,
 
     /// If present, the node's wireguard public key has changed.
-    pub key: Option<NodePublicKey>,
+    pub key: Option<PublicKey<ts_keys::Node>>,
 
     /// If present, the signature of the node's wireguard public key has changed.
     #[serde(borrow)]
     pub key_signature: Option<MarshaledSignature<'a>>,
 
     /// If present, the node's disco key has changed.
-    pub disco_key: Option<DiscoPublicKey>,
+    pub disco_key: Option<PublicKey<Disco>>,
     /// If present, the node's online status changed.
     pub online: Option<bool>,
     /// If present, the node's last seen time changed.

ts_control_serde/src/node.rs

@@ -4,7 +4,7 @@ use core::net::{IpAddr, SocketAddr};
 use chrono::{DateTime, Utc};
 use serde::{Deserialize, Serialize};
 use ts_capabilityversion::CapabilityVersion;
-use ts_keys::{DiscoPublicKey, MachinePublicKey, NodePublicKey};
+use ts_keys::{Disco, Machine, Node as NodeKey, PublicKey};
 
 use crate::{DnsResolver, derp_map::RegionId, host_info::HostInfo, user::UserId};
 
@@ -63,18 +63,18 @@ pub struct Node<'a> {
     /// Unique ID of the user who shared this node, if non-zero and different from [`Node::user`].
     pub sharer: UserId,
 
-    /// If populated, the public key of the Tailscale node's [`NodeKeyPair`][ts_keys::NodeKeyPair].
-    pub key: NodePublicKey,
-    /// The date and time that the Tailscale node's [`NodeKeyPair`][ts_keys::NodeKeyPair] will expire.
+    /// If populated,the Tailscale node's [node public key][PublicKey<ts_keys::Node>].
+    pub key: PublicKey<NodeKey>,
+    /// The date and time that the Tailscale node's node keypair will expire.
     pub key_expiry: Option<DateTime<Utc>>,
     /// If populated, a signature of the Tailnet Key Authority (TKA) key authorizing the Tailscale
     /// node to join the Tailnet.
     #[serde(borrow)]
     pub key_signature: Option<MarshaledSignature<'a>>,
-    /// If populated, the public key of the Tailscale node's [`MachineKeyPair`][ts_keys::MachineKeyPair].
-    pub machine: Option<MachinePublicKey>,
-    /// If populated, the public key of the Tailscale node's [`DiscoKeyPair`][ts_keys::DiscoKeyPair].
-    pub disco_key: Option<DiscoPublicKey>,
+    /// If populated, the Tailscale node's [machine public key][PublicKey<Machine>].
+    pub machine: Option<PublicKey<Machine>>,
+    /// If populated, the Tailscale node's [disco public key][PublicKey<Disco>].
+    pub disco_key: Option<PublicKey<Disco>>,
 
     /// The IP addresses of the Tailscale node in the Tailnet. There are exactly 2 addresses, and
     /// they are always in the same order: the first is the IPv4 address, the second is the IPv6

ts_control_serde/src/register.rs

@@ -2,7 +2,7 @@ use core::fmt::Debug;
 
 use chrono::{DateTime, Utc};
 use serde::{Deserialize, Serialize};
-use ts_keys::{NetworkLockPublicKey, NodePublicKey};
+use ts_keys::{NetworkLock, PublicKey};
 use url::Url;
 
 use crate::{
@@ -77,14 +77,14 @@ pub struct RegisterRequest<'a> {
     /// The current public key of this Tailscale node. In the case of node key rotation, this is
     /// the "new" node public key, and [`RegisterRequest::old_node_key`] contains the expired
     /// public node key.
-    pub node_key: NodePublicKey,
+    pub node_key: PublicKey<ts_keys::Node>,
     /// The expired public key of this Tailscale node. Only populated when the node key has expired
     /// and needs to be rotated.
-    pub old_node_key: Option<NodePublicKey>,
+    pub old_node_key: Option<PublicKey<ts_keys::Node>>,
     /// The new Tailnet Lock public key for this Tailscale node. Only populated when the key has
     /// been changed, or has never been set for this node.
     #[serde(rename = "NLKey")]
-    pub nl_key: Option<NetworkLockPublicKey>,
+    pub nl_key: Option<PublicKey<NetworkLock>>,
     /// Authentication information that allows this Tailscale node to register with the control
     /// plane and join a specific Tailnet.
     #[serde(borrow)]

ts_dataplane/Cargo.toml

@@ -25,6 +25,7 @@ ts_underlay_router.workspace = true
 ts_tunnel.workspace = true
 ts_bart.workspace = true
 ts_disco_protocol.workspace = true
+ts_keys.workspace = true
 
 # Unconditionally required dependencies.
 bytes.workspace = true

ts_dataplane/src/async_tokio.rs

@@ -3,9 +3,9 @@
 use std::{collections::HashMap, convert::Infallible, ops::DerefMut, sync::atomic::AtomicU32};
 
 use tokio::sync::{Mutex, mpsc};
+use ts_keys::{KeyPair, Node};
 use ts_packet::PacketMut;
 use ts_transport::{OverlayTransportId, PeerId, UnderlayTransportId};
-use ts_tunnel::NodeKeyPair;
 
 use crate::{EventResult, InboundResult, OutboundResult};
 
@@ -86,7 +86,7 @@ impl DataPlane {
     ///
     /// The second and third elements of the return tuple are output queues for disco and
     /// STUN messages, respectively.
-    pub fn new(my_key: NodeKeyPair) -> (Self, Rx<DiscoBatch>, Rx<StunBatch>) {
+    pub fn new(my_key: KeyPair<Node>) -> (Self, Rx<DiscoBatch>, Rx<StunBatch>) {
         let (overlay_up, overlay_down) = mpsc::unbounded_channel();
         let (underlay_down, underlay_up) = mpsc::unbounded_channel();
 

ts_dataplane/src/lib.rs

@@ -10,13 +10,14 @@ use ts_packet::PacketMut;
 use ts_packetfilter::{FilterExt, IpProto};
 use ts_time::{Handle, Scheduler};
 use ts_transport::{OverlayTransportId, PeerId, UnderlayTransportId};
-use ts_tunnel::{Endpoint, NodeKeyPair};
+use ts_tunnel::Endpoint;
 use ts_underlay_router as ur;
 
 pub mod async_tokio;
 mod packet_ident;
 
 pub use packet_ident::{PacketIdent, PacketType};
+use ts_keys::{KeyPair, Node};
 
 /// A data plane subsystem that can be the subject of timer events.
 pub enum Subsystem {
@@ -51,7 +52,7 @@ pub struct DataPlane {
 
 impl DataPlane {
     /// Creates a new data plane for a wireguard node key.
-    pub fn new(my_key: NodeKeyPair) -> Self {
+    pub fn new(my_key: KeyPair<Node>) -> Self {
         DataPlane {
             wireguard: Endpoint::new(my_key),
             or_out: Default::default(),

ts_derp/examples/listen.rs

@@ -2,7 +2,7 @@
 //!
 //! Intended to test ping/pong/keepalive.
 
-use ts_keys::NodeKeyPair;
+use ts_keys::KeyPair;
 
 mod common;
 
@@ -13,7 +13,7 @@ async fn main() -> ts_cli_util::Result<()> {
     let derp_map = common::load_derp_map().await;
     let region = derp_map.get(&common::REGION_1).unwrap();
 
-    let keypair = NodeKeyPair::new();
+    let keypair = KeyPair::random();
 
     let client = ts_derp::Client::connect(region, &keypair).await?;
     tracing::info!("derp handshake done");