Implement Agent Forwarding (#1249)

This PR implements Agent forwarding as discussed at
https://github.com/warp-tech/warpgate/issues/177#issuecomment-2647526617

Huge thanks to @Eugeny for the pointing me in the right direction

---------

Signed-off-by: Sam Toxopeus <sam@toxopeus.it>

README.md

@@ -116,6 +116,7 @@ Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/d
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/Gurkengewuerz"><img src="https://avatars.githubusercontent.com/u/10966337?v=4?s=100" width="100px;" alt="Niklas"/><br /><sub><b>Niklas</b></sub></a><br /><a href="https://github.com/Eugeny/warpgate/commits?author=Gurkengewuerz" title="Code">💻</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/notnooblord"><img src="https://avatars.githubusercontent.com/u/11678665?v=4?s=100" width="100px;" alt="Nooblord"/><br /><sub><b>Nooblord</b></sub></a><br /><a href="https://github.com/Eugeny/warpgate/commits?author=notnooblord" title="Code">💻</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://shea.nz/"><img src="https://avatars.githubusercontent.com/u/51303984?v=4?s=100" width="100px;" alt="Shea Smith"/><br /><sub><b>Shea Smith</b></sub></a><br /><a href="https://github.com/Eugeny/warpgate/commits?author=SheaSmith" title="Code">💻</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/samtoxie"><img src="https://avatars.githubusercontent.com/u/7732658?v=4?s=100" width="100px;" alt="samtoxie"/><br /><sub><b>samtoxie</b></sub></a><br /><a href="https://github.com/Eugeny/warpgate/commits?author=samtoxie" title="Code">💻</a></td>
     </tr>
   </tbody>
 </table>

warpgate-protocol-ssh/src/client/channel_session.rs

@@ -96,6 +96,11 @@ impl SessionChannel {
                                 request.x11_auth_cookie,
                                 request.x11_screen_number,
                             ).await?;
+                        },
+                        Some(ChannelOperation::AgentForward) => {
+                            self.client_channel.agent_forward(
+                                true,
+                            ).await?;
                         }
                         Some(ChannelOperation::Close) => break,
                         None => break,

warpgate-protocol-ssh/src/client/handler.rs

@@ -16,6 +16,7 @@ pub enum ClientHandlerEvent {
     HostKeyUnknown(PublicKey, oneshot::Sender<bool>),
     ForwardedTcpIp(Channel<Msg>, ForwardedTcpIpParams),
     ForwardedStreamlocal(Channel<Msg>, ForwardedStreamlocalParams),
+    ForwardedAgent(Channel<Msg>),
     X11(Channel<Msg>, String, u32),
     Disconnect,
 }
@@ -161,6 +162,17 @@ impl russh::client::Handler for ClientHandler {
         ));
         Ok(())
     }
+
+    async fn server_channel_open_agent_forward(
+        &mut self,
+        channel: Channel<Msg>,
+        _session: &mut Session,
+    ) -> Result<(), Self::Error> {
+        let _ = self.event_tx.send(ClientHandlerEvent::ForwardedAgent(
+            channel,
+        ));
+        Ok(())
+    }
 }
 
 impl Drop for ClientHandler {

warpgate-protocol-ssh/src/client/mod.rs

@@ -92,6 +92,7 @@ pub enum RCEvent {
     HostKeyUnknown(PublicKey, oneshot::Sender<bool>),
     ForwardedTcpIp(Uuid, ForwardedTcpIpParams),
     ForwardedStreamlocal(Uuid, ForwardedStreamlocalParams),
+    ForwardedAgent(Uuid),
     X11(Uuid, String, u32),
 }
 
@@ -319,6 +320,11 @@ impl RemoteClient {
                         let id = self.setup_server_initiated_channel(channel).await?;
                         let _ = self.tx.send(RCEvent::ForwardedStreamlocal(id, params));
                     }
+                    ClientHandlerEvent::ForwardedAgent(channel) => {
+                        info!("New forwarded agent connection");
+                        let id = self.setup_server_initiated_channel(channel).await?;
+                        let _ = self.tx.send(RCEvent::ForwardedAgent(id));
+                    }
                     ClientHandlerEvent::X11(channel, originator_address, originator_port) => {
                         info!("New X11 connection from {originator_address}:{originator_port:?}");
                         let id = self.setup_server_initiated_channel(channel).await?;

warpgate-protocol-ssh/src/common.rs

@@ -62,6 +62,7 @@ pub enum ChannelOperation {
     RequestEnv(String, String),
     RequestExec(String),
     RequestX11(X11Request),
+    AgentForward,
     RequestSubsystem(String),
     Data(Bytes),
     ExtendedData { data: Bytes, ext: u32 },

warpgate-protocol-ssh/src/server/russh_handler.rs

@@ -1,5 +1,4 @@
 use std::fmt::Debug;
-
 use bytes::Bytes;
 use russh::keys::PublicKey;
 use russh::server::{Auth, Handle, Msg, Session};
@@ -49,6 +48,7 @@ pub enum ServerHandlerEvent {
     CancelTcpIpForward(String, u32, oneshot::Sender<bool>),
     StreamlocalForward(String, oneshot::Sender<bool>),
     CancelStreamlocalForward(String, oneshot::Sender<bool>),
+    AgentForward(ServerChannelId, oneshot::Sender<bool>),
     Disconnect,
 }
 
@@ -516,6 +516,18 @@ impl russh::server::Handler for ServerHandler {
         }
         Ok(allowed)
     }
+
+    async fn agent_request(&mut self, channel: ChannelId, session: &mut Session) -> Result<bool, Self::Error> {
+        let (tx, rx) = oneshot::channel();
+        self.send_event(ServerHandlerEvent::AgentForward(ServerChannelId(channel), tx))?;
+        let allowed = rx.await.unwrap_or(false);
+        if allowed {
+            session.request_success()
+        } else {
+            session.request_failure()
+        }
+        Ok(allowed)
+    }
 }
 
 impl Drop for ServerHandler {

warpgate-protocol-ssh/src/server/session.rs

@@ -580,6 +580,12 @@ impl ServerSession {
                 self._cancel_streamlocal_forward(socket_path).await?;
                 let _ = reply.send(true);
             }
+
+            ServerHandlerEvent::AgentForward(channel,reply) => {
+                self._agent_forward(channel).await?;
+                let _ = reply.send(true);
+            }
+
             ServerHandlerEvent::Disconnect => (),
         }
 
@@ -865,6 +871,17 @@ impl ServerSession {
                     }
                 }
             }
+            RCEvent::ForwardedAgent(id) => {
+                if let Some(session) = &mut self.session_handle {
+                    let server_channel = session
+                        .channel_open_agent()
+                        .await?;
+
+                    self.channel_map
+                        .insert(ServerChannelId(server_channel.id()), id);
+                    self.all_channels.push(id);
+                }
+            }
             RCEvent::X11(id, originator_address, originator_port) => {
                 if let Some(session) = &mut self.session_handle {
                     let server_channel = session
@@ -1246,6 +1263,18 @@ impl ServerSession {
             .map_err(anyhow::Error::from)
     }
 
+    async fn _agent_forward(&mut self, server_channel_id: ServerChannelId) -> Result<()> {
+        let channel_id = self.map_channel(&server_channel_id)?;
+        debug!(channel=%channel_id, "Requested Agent Forwarding");
+        let _ = self.maybe_connect_remote().await;
+        self.send_command_and_wait(RCCommand::Channel(
+            channel_id,
+            ChannelOperation::AgentForward,
+        ))
+            .await?;
+        Ok(())
+    }
+
     async fn _auth_publickey_offer(
         &mut self,
         ssh_username: Secret<String>,