From 41bf13d1a42fee648596ee3023ee53b0ccc3bcc7 Mon Sep 17 00:00:00 2001
From: Matt Johnston <matt@ucc.asn.au>
Date: Sat, 8 Apr 2023 14:13:11 +0800
Subject: [PATCH] Simplify pubkey auth pkok

We don't need OwnedSig state any more (we send straight away), simplify
packet creation.
---
 src/cliauth.rs | 56 +++++++++++++++++---------------------------------
 src/packets.rs | 28 ++++++++++++-------------
 2 files changed, 33 insertions(+), 51 deletions(-)

diff --git a/src/cliauth.rs b/src/cliauth.rs
index 9ac2435..61327c4 100644
--- a/src/cliauth.rs
+++ b/src/cliauth.rs
@@ -23,34 +23,39 @@ use kex::SessId;
 use auth::AuthType;
 
 // pub for packets::ParseContext
-pub enum Req {
+enum Req {
     Password(ResponseString),
     PubKey { key: SignKey },
 }
 
-pub(crate) enum AuthState {
+enum AuthState {
     Unstarted,
     MethodQuery,
-    Request { last_req: Req, sig: Option<OwnedSig> },
+    Request { last_req: Req },
     Idle,
 }
 
 impl Req {
     // Creates a packet from the current request
+    // parse_ctx.cli_auth_type will be updated
+    // sig is an optional signature for pubkey auth
     fn req_packet<'b>(
         &'b self,
         username: &'b str,
         parse_ctx: &mut ParseContext,
+        sig: Option<&'b OwnedSig>,
     ) -> Result<Packet<'b>> {
+
         let username = username.into();
         let p = match self {
             Req::PubKey { key, .. } => {
                 // already checked by make_pubkey_req()
                 parse_ctx.cli_auth_type = Some(AuthType::PubKey);
+                let method = AuthMethod::PubKey(MethodPubKey::new(key.pubkey(), sig)?);
                 packets::UserauthRequest {
                     username,
                     service: SSH_SERVICE_CONNECTION,
-                    method: key.pubkey().try_into()?,
+                    method,
                 }.into()
             }
             Req::Password(pw) => {
@@ -215,42 +220,19 @@ impl CliAuth {
         s: &mut TrafSend<'_, '_>,
         b: &mut impl CliBehaviour,
     ) -> Result<()> {
-        // We are only sending keys one at a time so they shouldn't
-        // get out of sync. In future we could change it to send
-        // multiple requests pipelined, though unsure of server
-        // acceptance of that.
         match &mut self.state {
-            // Some tricky logistics to create the signature in self.state
-            // using a packet borrowed from other parts of self.state
-            AuthState::Request { last_req, ref mut sig } => {
-                if sig.is_some() {
-                    return Err(Error::SSHProtoError);
-                }
-                if let Req::PubKey { key, .. } = last_req {
+            AuthState::Request { last_req } => {
+                if let Req::PubKey { ref key } = last_req {
                     if key.pubkey() != pkok.key.0 {
+                        trace!("Received pkok for a different key");
                         return Err(Error::SSHProtoError);
                     }
-                }
 
-                let mut p = last_req.req_packet(&self.username, parse_ctx)?;
-                let last_req = &last_req;
-                if let Req::PubKey { key, .. } = last_req {
-                    // Create the signature
+                    // Sign the packet without the signature
+                    let p = last_req.req_packet(&self.username, parse_ctx, None)?;
                     let new_sig = Self::auth_sig_msg(&key, sess_id, &p, b).await?;
-                    let rsig = sig.insert(new_sig);
-
-                    // Put it in the packet
-                    if let Packet::UserauthRequest(UserauthRequest {
-                        method:
-                            AuthMethod::PubKey(MethodPubKey {
-                                sig: ref mut psig, ..
-                            }),
-                        ..
-                    }) = p
-                    {
-                        let rsig = &*rsig;
-                        *psig = Some(Blob(rsig.into()))
-                    }
+                    let p = last_req.req_packet(&self.username, parse_ctx, Some(&new_sig))?;
+
                     s.send(p)?;
                     return Ok(());
                 }
@@ -276,7 +258,7 @@ impl CliAuth {
             while self.try_pubkey {
                 let req = self.make_pubkey_req(b).await;
                 if let Some(req) = req {
-                    self.state = AuthState::Request { last_req: req, sig: None };
+                    self.state = AuthState::Request { last_req: req };
                     break;
                 }
             }
@@ -288,14 +270,14 @@ impl CliAuth {
         {
             let req = self.make_password_req(b).await?;
             if let Some(req) = req {
-                self.state = AuthState::Request { last_req: req, sig: None };
+                self.state = AuthState::Request { last_req: req };
             } else {
                 self.try_password = false;
             }
         }
 
         if let AuthState::Request { last_req, .. } = &self.state {
-            let p = last_req.req_packet(&self.username, parse_ctx)?;
+            let p = last_req.req_packet(&self.username, parse_ctx, None)?;
             s.send(p)?;
         } else {
             return Err(Error::BehaviourError {
diff --git a/src/packets.rs b/src/packets.rs
index 42fe5d6..087abb2 100644
--- a/src/packets.rs
+++ b/src/packets.rs
@@ -167,20 +167,6 @@ pub enum AuthMethod<'a> {
     Unknown(Unknown<'a>),
 }
 
-impl<'a> TryFrom<PubKey<'a>> for AuthMethod<'a> {
-    type Error = Error;
-    fn try_from(pubkey: PubKey<'a>) -> Result<Self> {
-        let sig_algo =
-            Signature::sig_name_for_pubkey(&pubkey).trap()?;
-        Ok(AuthMethod::PubKey(MethodPubKey {
-            sig_algo,
-            pubkey: Blob(pubkey),
-            sig: None,
-        }))
-    }
-}
-
-
 #[derive(Debug, SSHEncode)]
 #[sshwire(no_variant_names)]
 pub enum Userauth60<'a> {
@@ -238,6 +224,20 @@ pub struct MethodPubKey<'a> {
     pub sig: Option<Blob<Signature<'a>>>,
 }
 
+impl<'a> MethodPubKey<'a> {
+    pub fn new(pubkey: PubKey<'a>, sig: Option<&'a OwnedSig>) -> Result<Self> {
+        let sig_algo =
+            Signature::sig_name_for_pubkey(&pubkey).trap()?;
+        let sig = sig.map(|s| Blob((s).into()));
+        Ok(MethodPubKey {
+            sig_algo,
+            pubkey: Blob(pubkey),
+            sig,
+        })
+
+    }
+}
+
 impl SSHEncode for MethodPubKey<'_> {
     fn enc(&self, s: &mut dyn SSHSink) -> WireResult<()> {
         // byte      SSH_MSG_USERAUTH_REQUEST
-- 
GitLab