make kex more pluggable

ssh-transport/README.md

@@ -11,3 +11,4 @@ Other relevant RFCs:
 - [RFC 5656 Elliptic Curve Algorithm Integration in the Secure Shell Transport Layer](https://datatracker.ietf.org/doc/html/rfc5656)
 - [RFC 6668 SHA-2 Data Integrity Verification for the Secure Shell (SSH) Transport Layer Protocol](https://datatracker.ietf.org/doc/html/rfc6668)
 - [RFC 8709 Ed25519 and Ed448 Public Key Algorithms for the Secure Shell (SSH) Protocol](https://datatracker.ietf.org/doc/html/rfc8709)
+- [RFC 8731 Secure Shell (SSH) Key Exchange Method Using Curve25519 and Curve448](https://datatracker.ietf.org/doc/html/rfc8731)

ssh-transport/src/keys.rs

@@ -5,16 +5,45 @@ use subtle::ConstantTimeEq;
 use crate::{
     client_error,
     packet::{EncryptedPacket, MsgKind, Packet, RawPacket},
-    Msg, Result,
+    Msg, Result, SshRng,
 };
 
 #[derive(Clone, Copy)]
 pub struct KexAlgorithm {
     pub name: &'static str,
+    pub exchange: fn(
+        client_public_key: &[u8],
+        random: &mut (dyn SshRng + Send + Sync),
+    ) -> Result<KexAlgorithmOutput>,
+}
+pub struct KexAlgorithmOutput {
+    /// K
+    pub shared_secret: Vec<u8>,
+    /// Q_S
+    pub server_public_key: Vec<u8>,
 }
 
+/// <https://datatracker.ietf.org/doc/html/rfc8731>
 pub const KEX_CURVE_25519_SHA256: KexAlgorithm = KexAlgorithm {
     name: "curve25519-sha256",
+    exchange: |client_public_key, rng| {
+        let secret = x25519_dalek::EphemeralSecret::random_from_rng(crate::SshRngRandAdapter(rng));
+        let server_public_key = x25519_dalek::PublicKey::from(&secret); // Q_S
+
+        let Ok(arr) = <[u8; 32]>::try_from(client_public_key) else {
+            return Err(crate::client_error!(
+                "invalid x25519 public key length, should be 32, was: {}",
+                client_public_key.len()
+            ));
+        };
+        let client_public_key = x25519_dalek::PublicKey::from(arr);
+        let shared_secret = secret.diffie_hellman(&client_public_key); // K
+
+        Ok(KexAlgorithmOutput {
+            server_public_key: server_public_key.as_bytes().to_vec(),
+            shared_secret: shared_secret.as_bytes().to_vec(),
+        })
+    },
 };
 
 pub struct AlgorithmNegotiation<T> {
@@ -48,7 +77,7 @@ pub(crate) trait Keys: Send + Sync + 'static {
     fn encrypt_packet_to_msg(&mut self, packet: Packet, packet_number: u64) -> Msg;
 
     fn additional_mac_len(&self) -> usize;
-    fn rekey(&mut self, h: [u8; 32], k: [u8; 32]) -> Result<(), ()>;
+    fn rekey(&mut self, h: [u8; 32], k: &[u8]) -> Result<(), ()>;
 }
 
 pub(crate) struct Plaintext;
@@ -63,18 +92,18 @@ impl Keys for Plaintext {
     fn additional_mac_len(&self) -> usize {
         0
     }
-    fn rekey(&mut self, _: [u8; 32], _: [u8; 32]) -> Result<(), ()> {
+    fn rekey(&mut self, _: [u8; 32], _: &[u8]) -> Result<(), ()> {
         Err(())
     }
 }
 
 impl Session {
-    pub(crate) fn new(h: [u8; 32], k: [u8; 32]) -> Self {
+    pub(crate) fn new(h: [u8; 32], k: &[u8]) -> Self {
         Self::from_keys(h, h, k)
     }
 
     /// <https://datatracker.ietf.org/doc/html/rfc4253#section-7.2>
-    fn from_keys(session_id: [u8; 32], h: [u8; 32], k: [u8; 32]) -> Self {
+    fn from_keys(session_id: [u8; 32], h: [u8; 32], k: &[u8]) -> Self {
         let encryption_key_client_to_server =
             SshChaCha20Poly1305::new(derive_key(k, h, "C", session_id));
         let encryption_key_server_to_client =
@@ -114,7 +143,7 @@ impl Keys for Session {
         poly1305::BLOCK_SIZE
     }
 
-    fn rekey(&mut self, h: [u8; 32], k: [u8; 32]) -> Result<(), ()> {
+    fn rekey(&mut self, h: [u8; 32], k: &[u8]) -> Result<(), ()> {
         *self = Self::from_keys(self.session_id, h, k);
         Ok(())
     }
@@ -122,7 +151,7 @@ impl Keys for Session {
 
 /// Derive a key from the shared secret K and exchange hash H.
 /// <https://datatracker.ietf.org/doc/html/rfc4253#section-7.2>
-fn derive_key(k: [u8; 32], h: [u8; 32], letter: &str, session_id: [u8; 32]) -> [u8; 64] {
+fn derive_key(k: &[u8], h: [u8; 32], letter: &str, session_id: [u8; 32]) -> [u8; 64] {
     let sha2len = sha2::Sha256::output_size();
     let mut output = [0; 64];
 
@@ -254,7 +283,7 @@ impl SshChaCha20Poly1305 {
         // Now, MAC the length || content, and push that to the end.
         let mac = poly1305::Poly1305::new(&poly1305_key.into()).compute_unpadded(&bytes);
 
-        bytes.extend_from_slice(&mac);
+        bytes.extend_from_slice(mac.as_slice());
 
         EncryptedPacket::from_encrypted_full_bytes(bytes)
     }

ssh-transport/src/lib.rs

@@ -8,14 +8,13 @@ use std::{collections::VecDeque, mem::take};
 use ed25519_dalek::ed25519::signature::Signer;
 use keys::AlgorithmNegotiation;
 use packet::{
-    DhKeyExchangeInitPacket, DhKeyExchangeInitReplyPacket, KeyExchangeInitPacket, Packet,
+    DhKeyExchangeInitReplyPacket, KeyExchangeEcDhInitPacket, KeyExchangeInitPacket, Packet,
     PacketTransport, SshPublicKey, SshSignature,
 };
-use parse::{MpInt, NameList, Parser, Writer};
+use parse::{NameList, Parser, Writer};
 use rand::RngCore;
 use sha2::Digest;
 use tracing::{debug, info, trace};
-use x25519_dalek::{EphemeralSecret, PublicKey};
 
 pub use packet::Msg;
 
@@ -64,10 +63,11 @@ enum ServerState {
         client_identification: Vec<u8>,
         client_kexinit: Vec<u8>,
         server_kexinit: Vec<u8>,
+        kex_algorithm: keys::KexAlgorithm,
     },
     NewKeys {
         h: [u8; 32],
-        k: [u8; 32],
+        k: Vec<u8>,
     },
     ServiceRequest,
     Open,
@@ -250,24 +250,24 @@ impl ServerConnection {
                         client_identification,
                         client_kexinit: packet.payload,
                         server_kexinit: server_kexinit_payload,
+                        kex_algorithm,
                     };
                 }
                 ServerState::DhKeyInit {
                     client_identification,
                     client_kexinit,
                     server_kexinit,
+                    kex_algorithm,
                 } => {
                     // TODO: move to keys.rs
-                    let dh = DhKeyExchangeInitPacket::parse(&packet.payload)?;
+                    let dh = KeyExchangeEcDhInitPacket::parse(&packet.payload)?;
 
-                    let secret =
+                    let client_public_key = dh.qc;
-                        EphemeralSecret::random_from_rng(SshRngRandAdapter(&mut *self.rng));
-                    let server_public_key = PublicKey::from(&secret); // Q_S
 
-                    let client_public_key = dh.e; // Q_C
+                    let keys::KexAlgorithmOutput {
-
+                        server_public_key,
-                    let shared_secret =
+                        shared_secret,
-                        secret.diffie_hellman(&client_public_key.as_x25519_public_key()?); // K
+                    } = (kex_algorithm.exchange)(client_public_key, &mut *self.rng)?;
 
                     let pub_hostkey = SshPublicKey {
                         format: b"ssh-ed25519",
@@ -297,12 +297,13 @@ impl ServerConnection {
                     hash_string(&mut hash, client_kexinit); // I_C
                     hash_string(&mut hash, server_kexinit); // I_S
                     add_hash(&mut hash, &pub_hostkey.to_bytes()); // K_S
+
                     // For normal DH as in RFC4253, e and f are mpints.
                     // But for ECDH as defined in RFC5656, Q_C and Q_S are strings.
                     // <https://datatracker.ietf.org/doc/html/rfc5656#section-4>
-                    hash_string(&mut hash, client_public_key.0); // Q_C
+                    hash_string(&mut hash, client_public_key); // Q_C
-                    hash_string(&mut hash, server_public_key.as_bytes()); // Q_S
+                    hash_string(&mut hash, &server_public_key); // Q_S
-                    hash_mpint(&mut hash, shared_secret.as_bytes()); // K
+                    hash_mpint(&mut hash, &shared_secret); // K
 
                     let hash = hash.finalize();
 
@@ -317,8 +318,8 @@ impl ServerConnection {
                     // eprintln!("hash:              {:x?}", hash);
 
                     let packet = DhKeyExchangeInitReplyPacket {
-                        pubkey: pub_hostkey,
+                        public_host_key: pub_hostkey,
-                        f: MpInt(server_public_key.as_bytes()),
+                        ephemeral_public_key: &server_public_key,
                         signature: SshSignature {
                             format: b"ssh-ed25519",
                             data: &signature.to_bytes(),
@@ -329,7 +330,7 @@ impl ServerConnection {
                     });
                     self.state = ServerState::NewKeys {
                         h: hash.into(),
-                        k: shared_secret.to_bytes(),
+                        k: shared_secret,
                     };
                 }
                 ServerState::NewKeys { h, k } => {
@@ -337,13 +338,11 @@ impl ServerConnection {
                         return Err(client_error!("did not send SSH_MSG_NEWKEYS"));
                     }
 
-                    let (h, k) = (*h, *k);
-
                     self.packet_transport.queue_packet(Packet {
                         payload: vec![Packet::SSH_MSG_NEWKEYS],
                     });
+                    self.packet_transport.set_key(*h, k);
                     self.state = ServerState::ServiceRequest {};
-                    self.packet_transport.set_key(h, k);
                 }
                 ServerState::ServiceRequest => {
                     // TODO: this should probably move out of here? unsure.

ssh-transport/src/packet.rs

@@ -4,7 +4,7 @@ use std::collections::VecDeque;
 
 use crate::client_error;
 use crate::keys::{Keys, Plaintext, Session};
-use crate::parse::{MpInt, NameList, Parser, Writer};
+use crate::parse::{NameList, Parser, Writer};
 use crate::Result;
 
 /// Frames the byte stream into packets.
@@ -101,7 +101,7 @@ impl PacketTransport {
         self.send_packets.pop_front()
     }
 
-    pub(crate) fn set_key(&mut self, h: [u8; 32], k: [u8; 32]) {
+    pub(crate) fn set_key(&mut self, h: [u8; 32], k: &[u8]) {
         if let Err(()) = self.keys.rekey(h, k) {
             self.keys = Box::new(Session::new(h, k));
         }
@@ -147,7 +147,9 @@ impl Packet {
 
     // 30 to 49 Key exchange method specific (numbers can be reused for different authentication methods)
     pub const SSH_MSG_KEXDH_INIT: u8 = 30;
+    pub const SSH_MSG_KEX_ECDH_INIT: u8 = 30; // Same number
     pub const SSH_MSG_KEXDH_REPLY: u8 = 31;
+    pub const SSH_MSG_KEX_ECDH_REPLY: u8 = 31;
 
     // -----
     // User authentication protocol:
@@ -329,21 +331,21 @@ impl<'a> KeyExchangeInitPacket<'a> {
 }
 
 #[derive(Debug)]
-pub(crate) struct DhKeyExchangeInitPacket<'a> {
+pub(crate) struct KeyExchangeEcDhInitPacket<'a> {
-    pub(crate) e: MpInt<'a>,
+    pub(crate) qc: &'a [u8],
 }
-impl<'a> DhKeyExchangeInitPacket<'a> {
+impl<'a> KeyExchangeEcDhInitPacket<'a> {
-    pub(crate) fn parse(payload: &'a [u8]) -> Result<DhKeyExchangeInitPacket<'_>> {
+    pub(crate) fn parse(payload: &'a [u8]) -> Result<KeyExchangeEcDhInitPacket<'_>> {
         let mut c = Parser::new(payload);
 
         let kind = c.u8()?;
-        if kind != Packet::SSH_MSG_KEXDH_INIT {
+        if kind != Packet::SSH_MSG_KEX_ECDH_INIT {
             return Err(client_error!(
                 "expected SSH_MSG_KEXDH_INIT packet, found {kind}"
             ));
         }
-        let e = c.mpint()?;
+        let qc = c.string()?;
-        Ok(Self { e })
+        Ok(Self { qc })
     }
 }
 
@@ -371,17 +373,19 @@ pub(crate) struct SshSignature<'a> {
 
 #[derive(Debug)]
 pub(crate) struct DhKeyExchangeInitReplyPacket<'a> {
-    pub(crate) pubkey: SshPublicKey<'a>,
+    /// K_S
-    pub(crate) f: MpInt<'a>,
+    pub(crate) public_host_key: SshPublicKey<'a>,
+    /// Q_S
+    pub(crate) ephemeral_public_key: &'a [u8],
     pub(crate) signature: SshSignature<'a>,
 }
 impl<'a> DhKeyExchangeInitReplyPacket<'a> {
     pub(crate) fn to_bytes(&self) -> Vec<u8> {
         let mut data = Writer::new();
 
-        data.u8(Packet::SSH_MSG_KEXDH_REPLY);
+        data.u8(Packet::SSH_MSG_KEX_ECDH_REPLY);
-        data.write(&self.pubkey.to_bytes());
+        data.write(&self.public_host_key.to_bytes());
-        data.mpint(self.f);
+        data.string(self.ephemeral_public_key);
 
         data.u32((4 + self.signature.format.len() + 4 + self.signature.data.len()) as u32);
         // <https://datatracker.ietf.org/doc/html/rfc8709#section-6>
@@ -480,7 +484,9 @@ impl PacketParser {
         // 'padding_length', 'payload', 'random padding', and 'mac').
         // Implementations SHOULD support longer packets, where they might be needed.
         if packet_length > 500_000 {
-            return Err(client_error!("packet too large (>500_000): {packet_length}"));
+            return Err(client_error!(
+                "packet too large (>500_000): {packet_length}"
+            ));
         }
 
         let remaining_len = std::cmp::min(bytes.len(), packet_length - (self.raw_data.len() - 4));

ssh-transport/src/parse.rs

@@ -58,8 +58,7 @@ impl<'a> Parser<'a> {
     }
 
     pub fn mpint(&mut self) -> Result<MpInt<'a>> {
-        let data = self.string()?;
+        todo!("do correctly")
-        Ok(MpInt(data))
     }
 
     pub fn string(&mut self) -> Result<&'a [u8]> {
@@ -101,8 +100,8 @@ impl Writer {
         self.string(list.0.as_bytes());
     }
 
-    pub fn mpint(&mut self, mpint: MpInt<'_>) {
+    pub fn mpint(&mut self, _mpint: MpInt<'_>) {
-        self.string(mpint.0);
+        todo!("implement correctly?")
     }
 
     pub fn string(&mut self, data: &[u8]) {
@@ -143,19 +142,5 @@ impl Debug for NameList<'_> {
     }
 }
 
-// TODO: THIS IS A BRITTLE MESS BECAUSE THE RFC SUCKS HERE
-// DO NOT TOUCH MPINT ENCODING ANYWHERE
 #[derive(Debug, Clone, Copy)]
-pub struct MpInt<'a>(pub(crate) &'a [u8]);
+pub struct MpInt<'a>(pub &'a [u8]);
-
-impl<'a> MpInt<'a> {
-    pub(crate) fn as_x25519_public_key(&self) -> Result<x25519_dalek::PublicKey> {
-        let Ok(arr) = <[u8; 32]>::try_from(self.0) else {
-            return Err(crate::client_error!(
-                "invalid x25519 public key length, should be 32, was: {}",
-                self.0.len()
-            ));
-        };
-        Ok(x25519_dalek::PublicKey::from(arr))
-    }
-}