Merge pull request #45 from wiktor-k/jcspencer/draft-14

Add SSH_AGENT_EXTENSION_RESPONSE message type (RFC Draft 14)

Author: wiktor-k

examples/key_storage.rs

@@ -11,7 +11,7 @@ use sha1::Sha1;
 use ssh_agent_lib::agent::NamedPipeListener as Listener;
 use ssh_agent_lib::agent::Session;
 use ssh_agent_lib::error::AgentError;
-use ssh_agent_lib::proto::extension::SessionBind;
+use ssh_agent_lib::proto::extension::{QueryResponse, RestrictDestination, SessionBind};
 use ssh_agent_lib::proto::{
     message, signature, AddIdentity, AddIdentityConstrained, AddSmartcardKeyConstrained,
     Credential, Extension, KeyConstraint, RemoveIdentity, SignRequest, SmartcardKey,
@@ -151,12 +151,13 @@ impl Session for KeyStorage {
         } = identity;
         info!("Would use these constraints: {constraints:#?}");
         for constraint in constraints {
-            if let KeyConstraint::Extension(name, mut details) = constraint {
-                if name == "restrict-destination-v00@openssh.com" {
-                    if let Ok(destination_constraint) = details.parse::<SessionBind>() {
-                        info!("Destination constraint: {destination_constraint:?}");
-                    }
+            if let KeyConstraint::Extension(mut extension) = constraint {
+                if let Some(destination) =
+                    extension.parse_key_constraint::<RestrictDestination>()?
+                {
+                    info!("Destination constraint: {destination:?}");
                 }
+
                 if let Credential::Key { privkey, comment } = identity.credential.clone() {
                     let privkey = PrivateKey::try_from(privkey).map_err(AgentError::other)?;
                     self.identity_add(Identity {
@@ -199,13 +200,28 @@ impl Session for KeyStorage {
         Ok(())
     }
 
-    async fn extension(&mut self, mut extension: Extension) -> Result<(), AgentError> {
+    async fn extension(
+        &mut self,
+        mut extension: Extension,
+    ) -> Result<Option<Extension>, AgentError> {
         info!("Extension: {extension:?}");
-        if extension.name == "session-bind@openssh.com" {
-            let bind = extension.details.parse::<SessionBind>()?;
-            info!("Bind: {bind:?}");
+
+        match extension.name.as_str() {
+            "query" => {
+                let response = Extension::new_message(QueryResponse {
+                    extensions: vec!["query".into(), "session-bind@openssh.com".into()],
+                })?;
+                Ok(Some(response))
+            }
+            "session-bind@openssh.com" => match extension.parse_message::<SessionBind>()? {
+                Some(bind) => {
+                    info!("Bind: {bind:?}");
+                    Ok(None)
+                }
+                None => Err(AgentError::Failure),
+            },
+            _ => Err(AgentError::Failure),
         }
-        Ok(())
     }
 }
 

src/agent.rs

@@ -176,7 +176,7 @@ pub trait Session: 'static + Sync + Send + Sized {
     }
 
     /// Invoke a custom, vendor-specific extension on the agent.
-    async fn extension(&mut self, _extension: Extension) -> Result<(), AgentError> {
+    async fn extension(&mut self, _extension: Extension) -> Result<Option<Extension>, AgentError> {
         Err(AgentError::from(ProtoError::UnsupportedCommand {
             command: 27,
         }))
@@ -206,7 +206,12 @@ pub trait Session: 'static + Sync + Send + Sized {
             Request::AddSmartcardKeyConstrained(key) => {
                 self.add_smartcard_key_constrained(key).await?
             }
-            Request::Extension(extension) => self.extension(extension).await?,
+            Request::Extension(extension) => {
+                return match self.extension(extension).await? {
+                    Some(response) => Ok(Response::ExtensionResponse(response)),
+                    None => Ok(Response::Success),
+                }
+            }
         }
         Ok(Response::Success)
     }
@@ -224,6 +229,10 @@ pub trait Session: 'static + Sync + Send + Sized {
                 log::debug!("Request: {incoming_message:?}");
                 let response = match self.handle(incoming_message).await {
                     Ok(message) => message,
+                    Err(AgentError::ExtensionFailure) => {
+                        log::error!("Extension failure handling message");
+                        Response::ExtensionFailure
+                    }
                     Err(e) => {
                         log::error!("Error handling message: {:?}", e);
                         Response::Failure

src/error.rs

@@ -20,6 +20,14 @@ pub enum AgentError {
     /// Other unspecified error.
     #[error("Other error: {0:#}")]
     Other(#[from] Box<dyn std::error::Error + Send + Sync + 'static>),
+
+    /// Generic agent extension failure
+    #[error("Generic agent extension failure")]
+    ExtensionFailure,
+
+    /// Generic agent failure
+    #[error("Generic agent failure")]
+    Failure,
 }
 
 impl AgentError {