ts_keys: rework to use generic types with marker traits

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

Author: danderson

Cargo.lock

@@ -336,9 +336,8 @@ 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,
+        DiscoKey, MachineKey, NetworkLockKey, NodeKey, NodeState, Pair, PersistState, Private,
+        Public,
     };
 }
 

src/lib.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::{MachineKey, Pair, Public};
 use url::Url;
 use zerocopy::network_endian::U32;
 
@@ -140,8 +140,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: Public<MachineKey>,
+    public_key: Public<MachineKey>,
 }
 
 impl fmt::Display for ControlPublicKeys {
@@ -153,7 +153,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: &Pair<MachineKey>,
 ) -> Result<Http2<BytesBody>, ConnectionError> {
     let h1_client = connect_h1(control_url).await?;
 
@@ -189,7 +189,7 @@ async fn connect_h1(url: &Url) -> Result<ts_http_util::Http1<EmptyBody>, Connect
 }
 
 #[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<Public<MachineKey>, ConnectionError> {
     let mut key_url = control_url.join("/key")?;
 
     #[cfg(not(feature = "insecure-keyfetch"))]
@@ -226,7 +226,7 @@ pub async fn upgrade_ts2021(
     control_url: &Url,
     init_msg: &str,
     handshake: ts_control_noise::Handshake,
-    machine_key: &MachineKeyPair,
+    machine_key: &Pair<MachineKey>,
     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/client/connect.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::{MachineKey, Pair};
 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: &Pair<MachineKey>,
     ) -> 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: &Pair<MachineKey>,
     stream: Io,
 ) -> Result<Http2<BytesBody>, Error>
 where

ts_control/src/control_dialer.rs

@@ -6,7 +6,7 @@ use core::{
 
 use chrono::{DateTime, Utc};
 use ts_capabilityversion::CapabilityVersion;
-use ts_keys::{DiscoPublicKey, MachinePublicKey, NodePublicKey};
+use ts_keys::{DiscoKey, MachineKey, NodeKey, Public};
 
 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: Public<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<Public<MachineKey>>,
+    /// The node's disco public key, if known.
+    pub disco_key: Option<Public<DiscoKey>>,
     /// 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<Public<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<Public<DiscoKey>>,
 
     /// The node's key signature for Tailnet Lock. If `None`, has not changed.
     pub tailnet_lock_key_signature: Option<Vec<u8>>,

ts_control/src/node.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::{MachineKey, Pair, Public};
 use ts_noise::ik::SentHandshake;
 use zerocopy::{IntoBytes, TryFromBytes};
 
@@ -31,14 +31,14 @@ 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: &Pair<MachineKey>,
+        control_public_key: &Public<MachineKey>,
         capability_version: ts_capabilityversion::CapabilityVersion,
     ) -> (Self, String) {
         let mut ciphertext = [0; SentHandshake::INIT_SIZE];
         let state = SentHandshake::new(
             node_machine_key.into(),
-            control_public_key.into(),
+            &control_public_key.into(),
             prologue.as_bytes(),
             &mut 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: &Pair<MachineKey>,
     ) -> Result<WrappedIo<T>, Error> {
         let mut hdr_bytes = [0u8; 3];
         conn.read_exact(&mut hdr_bytes[..]).await?;
@@ -78,7 +78,7 @@ 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.into()) {
             Ok(session) => session,
             Err(state) => {
                 self.state = state;

ts_control_noise/src/handshake.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::{DiscoKey, Public};
 
 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: Public<ts_keys::NodeKey>,
     /// 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: Public<DiscoKey>,
 
     /// 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<Public<ts_keys::NodeKey>>,
 
     /// 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<Public<DiscoKey>>,
     /// 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/netmap.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::{DiscoKey, MachineKey, NodeKey, Public};
 
 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 public key of the Tailscale node's [node keypair][ts_keys::Pair<ts_keys::NodeKey>].
+    pub key: Public<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 public key of the Tailscale node's [machine keypair][ts_keys::Pair<ts_keys::MachineKey>].
+    pub machine: Option<Public<MachineKey>>,
+    /// If populated, the public key of the Tailscale node's [disco keypair][ts_keys::Pair<DiscoKey>].
+    pub disco_key: Option<Public<DiscoKey>>,
 
     /// 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/node.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::{NetworkLockKey, Public};
 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: Public<ts_keys::NodeKey>,
     /// 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<Public<ts_keys::NodeKey>>,
     /// 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<Public<NetworkLockKey>>,
     /// Authentication information that allows this Tailscale node to register with the control
     /// plane and join a specific Tailnet.
     #[serde(borrow)]

ts_control_serde/src/register.rs

@@ -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