From a515a9e7385748cc25f7ffd99623ee0f7f0ad3d8 Mon Sep 17 00:00:00 2001
From: George Hopkins <george-hopkins@null.net>
Date: Wed, 1 Nov 2023 09:53:40 +0100
Subject: [PATCH] Fix handling of key constraints

---
 russh-keys/src/agent/client.rs | 18 +++++++++++++-----
 russh-keys/src/agent/msg.rs    |  3 ++-
 russh-keys/src/agent/server.rs |  6 ++----
 3 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/russh-keys/src/agent/client.rs b/russh-keys/src/agent/client.rs
index 1e730030..23ac90e3 100644
--- a/russh-keys/src/agent/client.rs
+++ b/russh-keys/src/agent/client.rs
@@ -87,6 +87,15 @@ impl<S: AsyncRead + AsyncWrite + Unpin> AgentClient<S> {
         Ok(())
     }
 
+    async fn read_success(&mut self) -> Result<(), Error> {
+        self.read_response().await?;
+        if self.buf.first() == Some(&msg::SUCCESS) {
+            Ok(())
+        } else {
+            Err(Error::AgentFailure)
+        }
+    }
+
     /// Send a key to the agent, with a (possibly empty) slice of
     /// constraints to apply when using the key to sign.
     pub async fn add_identity(
@@ -131,12 +140,11 @@ impl<S: AsyncRead + AsyncWrite + Unpin> AgentClient<S> {
             }
         }
         if !constraints.is_empty() {
-            self.buf.push_u32_be(constraints.len() as u32);
             for cons in constraints {
                 match *cons {
                     Constraint::KeyLifetime { seconds } => {
                         self.buf.push(msg::CONSTRAIN_LIFETIME);
-                        self.buf.push_u32_be(seconds)
+                        self.buf.push_u32_be(seconds);
                     }
                     Constraint::Confirm => self.buf.push(msg::CONSTRAIN_CONFIRM),
                     Constraint::Extensions {
@@ -153,7 +161,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> AgentClient<S> {
         let len = self.buf.len() - 4;
         BigEndian::write_u32(&mut self.buf[..], len as u32);
 
-        self.read_response().await?;
+        self.read_success().await?;
         Ok(())
     }
 
@@ -467,8 +475,8 @@ impl<S: AsyncRead + AsyncWrite + Unpin> AgentClient<S> {
         self.buf.clear();
         self.buf.resize(4);
         self.buf.push(msg::REMOVE_ALL_IDENTITIES);
-        BigEndian::write_u32(&mut self.buf[..], 5);
-        self.read_response().await?;
+        BigEndian::write_u32(&mut self.buf[..], 1);
+        self.read_success().await?;
         Ok(())
     }
 
diff --git a/russh-keys/src/agent/msg.rs b/russh-keys/src/agent/msg.rs
index a77c5091..d732e674 100644
--- a/russh-keys/src/agent/msg.rs
+++ b/russh-keys/src/agent/msg.rs
@@ -19,4 +19,5 @@ pub const EXTENSION: u8 = 27;
 
 pub const CONSTRAIN_LIFETIME: u8 = 1;
 pub const CONSTRAIN_CONFIRM: u8 = 2;
-pub const CONSTRAIN_EXTENSION: u8 = 3;
+// pub const CONSTRAIN_MAXSIGN: u8 = 3;
+pub const CONSTRAIN_EXTENSION: u8 = 255;
diff --git a/russh-keys/src/agent/server.rs b/russh-keys/src/agent/server.rs
index 7fc2450c..af3891f3 100644
--- a/russh-keys/src/agent/server.rs
+++ b/russh-keys/src/agent/server.rs
@@ -325,10 +325,8 @@ impl<S: AsyncRead + AsyncWrite + Send + Unpin + 'static, A: Agent + Send + Sync
         let mut w = self.keys.0.write().or(Err(Error::AgentFailure))?;
         let now = SystemTime::now();
         if constrained {
-            let n = r.read_u32()?;
             let mut c = Vec::new();
-            for _ in 0..n {
-                let t = r.read_byte()?;
+            while let Ok(t) = r.read_byte() {
                 if t == msg::CONSTRAIN_LIFETIME {
                     let seconds = r.read_u32()?;
                     c.push(Constraint::KeyLifetime { seconds });
@@ -353,7 +351,7 @@ impl<S: AsyncRead + AsyncWrite + Send + Unpin + 'static, A: Agent + Send + Sync
                     return Ok(false);
                 }
             }
-            w.insert(blob, (Arc::new(key), now, Vec::new()));
+            w.insert(blob, (Arc::new(key), now, c));
         } else {
             w.insert(blob, (Arc::new(key), now, Vec::new()));
         }