From 8bd4025ba3b8da8875f9db7db9a884accca742d7 Mon Sep 17 00:00:00 2001
From: Thanh Nguyen <thanh@timweri.me>
Date: Tue, 28 Nov 2023 18:51:34 -0500
Subject: [PATCH] Handle pubkey blob in sign call being SSH cert

---
 rustica-agent/src/lib.rs | 22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/rustica-agent/src/lib.rs b/rustica-agent/src/lib.rs
index 8c84ba7..c392c4a 100644
--- a/rustica-agent/src/lib.rs
+++ b/rustica-agent/src/lib.rs
@@ -291,6 +291,7 @@ impl SshAgentHandler for Handler {
         _flags: u32,
     ) -> Result<Response, AgentError> {
         trace!("Sign call");
+
         // Tri check to find how to sign the request. Since starting rustica with a file based
         // key is the same process as keys added afterwards, we do this to prevent duplication
         // of the private key based signing code.
@@ -333,13 +334,27 @@ impl SshAgentHandler for Handler {
 
             return Ok(Response::SignResponse { signature });
         } else if let Signatory::Direct(privkey) = &self.signatory {
+            // Extract the pubkey fingerprint from either the SSH pubkey or the SSH cert
+            let fingerprint = match (Certificate::from_bytes(&pubkey), PublicKey::from_bytes(&pubkey)) {
+                (Ok(cert), _) => cert.key.fingerprint(),
+                (_, Ok(pubkey)) => pubkey.fingerprint(),
+                _ => return Err(AgentError::from("Invalid key blob")),
+            };
+
             // Don't sign requests if the requested key does not match the signatory
-            if privkey.pubkey.encode() != pubkey {
+            if privkey.pubkey.fingerprint() != fingerprint {
                 return Err(AgentError::from("No such key"));
             }
 
             Some(privkey)
         } else if let Signatory::Yubikey(signer) = &mut self.signatory {
+            // Extract the pubkey fingerprint from either the SSH pubkey or the SSH cert
+            let fingerprint = match (Certificate::from_bytes(&pubkey), PublicKey::from_bytes(&pubkey)) {
+                (Ok(cert), _) => cert.key.fingerprint(),
+                (_, Ok(pubkey)) => pubkey.fingerprint(),
+                _ => return Err(AgentError::from("Invalid key blob")),
+            };
+
             // Don't sign requests if the requested key does not match the signatory
             if signer
                 .yk
@@ -348,11 +363,12 @@ impl SshAgentHandler for Handler {
                     println!("Yubikey Fetch Certificate Error: {e}");
                     AgentError::from("Yubikey fetch certificate error")
                 })?
-                .encode()
-                != pubkey
+                .fingerprint()
+                != fingerprint
             {
                 return Err(AgentError::from("No such key"));
             }
+
             // Since we are using the Yubikey for a signing operation the only time they
             // won't have to tap here is if they are using cached keys and this is right after
             // a secure Rustica tap. In most cases, we'll need to send this, rarely, it'll be