Skip to content

Commit

Permalink
Merge pull request #65 from wiktor-k/wiktor/add-proto-dumper
Browse files Browse the repository at this point in the history
Add protocol dumper example
  • Loading branch information
wiktor-k authored May 13, 2024
2 parents 73fcf48 + 181d7ce commit 1b556c8
Show file tree
Hide file tree
Showing 4 changed files with 137 additions and 13 deletions.
10 changes: 5 additions & 5 deletions examples/pgp-wrapper.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
//!
//! Works perfectly in conjunction with `openpgp-card-agent.rs`!
use std::{cell::RefCell, pin::Pin};
use std::cell::RefCell;

use chrono::DateTime;
use clap::Parser;
Expand All @@ -55,11 +55,11 @@ use tokio::runtime::Runtime;
struct WrappedKey {
public_key: PublicKey,
pubkey: KeyData,
client: RefCell<Pin<Box<dyn Session>>>,
client: RefCell<Box<dyn Session>>,
}

impl WrappedKey {
fn new(pubkey: KeyData, client: Pin<Box<dyn Session>>) -> Self {
fn new(pubkey: KeyData, client: Box<dyn Session>) -> Self {
let KeyData::Ed25519(key) = pubkey.clone() else {
panic!("The first key was not ed25519!");
};
Expand Down Expand Up @@ -206,11 +206,11 @@ fn main() -> testresult::TestResult {
let (client, identities) = rt.block_on(async move {
#[cfg(unix)]
let mut client =
connect(Binding::FilePath(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?).await?;
connect(Binding::FilePath(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?)?;

#[cfg(windows)]
let mut client =
connect(Binding::NamedPipe(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?).await?;
connect(Binding::NamedPipe(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?)?;

let identities = client.request_identities().await?;

Expand Down
124 changes: 124 additions & 0 deletions examples/proto-dumper.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,124 @@
//! This example illustrates a couple of features: First, it
//! implements a forwarder, exposing an SSH agent socket and
//! forwarding to a different one. Secondly it shows how to work with
//! low-level handling of messages instead of parsed high-level
//! structures.
//!
//! Run with
//! RUST_LOG=info cargo run --example proto-dumper -- --target unix://$SSH_AUTH_SOCK -H unix:///tmp/test.sock
use clap::Parser;
use service_binding::Binding;
use ssh_agent_lib::{
agent::bind,
agent::Agent,
agent::Session,
async_trait,
client::connect,
error::AgentError,
proto::{Request, Response},
};
use ssh_encoding::Encode;

struct DumpAndForward {
target: Box<dyn Session>,
session: u64,
id: u64,
}

#[async_trait]
impl Session for DumpAndForward {
async fn handle(&mut self, message: Request) -> Result<Response, AgentError> {
use std::io::Write;

self.id += 1;
let req_file = format!("req-{}-{}.bin", self.session, self.id);
log::info!("Writing request {message:?} to {req_file}");

let mut req = std::fs::File::create(req_file)?;
let mut buf = vec![];
message.encode(&mut buf).map_err(AgentError::other)?;
req.write_all(&buf)?;
drop(req);

let response = self.target.handle(message).await?;

let resp_file = format!("resp-{}-{}.bin", self.session, self.id);
log::info!("Writing response {response:?} to {resp_file}");
let mut resp = std::fs::File::create(resp_file)?;
let mut buf = vec![];
response.encode(&mut buf).map_err(AgentError::other)?;
resp.write_all(&buf)?;
drop(resp);

Ok(response)
}
}

struct Forwarder {
target: Binding,
id: u64,
}

#[cfg(unix)]
impl Agent<tokio::net::UnixListener> for Forwarder {
fn new_session(&mut self, _socket: &tokio::net::UnixStream) -> impl Session {
self.create_new_session()
}
}

impl Agent<tokio::net::TcpListener> for Forwarder {
fn new_session(&mut self, _socket: &tokio::net::TcpStream) -> impl Session {
self.create_new_session()
}
}

#[cfg(windows)]
impl Agent<ssh_agent_lib::agent::NamedPipeListener> for Forwarder {
fn new_session(
&mut self,
_socket: &tokio::net::windows::named_pipe::NamedPipeServer,
) -> impl Session {
self.create_new_session()
}
}

impl Forwarder {
fn create_new_session(&mut self) -> impl Session {
self.id += 1;
DumpAndForward {
target: connect(self.target.clone().try_into().unwrap()).unwrap(),
session: self.id,
id: 0,
}
}
}

#[derive(Debug, Parser)]
struct Args {
/// Target SSH agent to which we will proxy all requests.
#[clap(long)]
target: Binding,

/// Source that we will bind to.
#[clap(long, short = 'H')]
host: Binding,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
env_logger::init();

let args = Args::parse();

bind(
args.host.try_into()?,
Forwarder {
target: args.target,
id: 0,
},
)
.await?;

Ok(())
}
4 changes: 2 additions & 2 deletions examples/ssh-agent-client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,11 @@ use ssh_agent_lib::client::connect;
async fn main() -> Result<(), Box<dyn std::error::Error>> {
#[cfg(unix)]
let mut client =
connect(Binding::FilePath(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?).await?;
connect(Binding::FilePath(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?)?;

#[cfg(windows)]
let mut client =
connect(Binding::NamedPipe(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?).await?;
connect(Binding::NamedPipe(std::env::var("SSH_AUTH_SOCK")?.into()).try_into()?)?;

eprintln!(
"Identities that this agent knows of: {:#?}",
Expand Down
12 changes: 6 additions & 6 deletions src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,18 +37,18 @@ where
}

/// Wrap a stream into an SSH agent client.
pub async fn connect(
pub fn connect(
stream: service_binding::Stream,
) -> Result<std::pin::Pin<Box<dyn crate::agent::Session>>, Box<dyn std::error::Error>> {
) -> Result<Box<dyn crate::agent::Session>, Box<dyn std::error::Error>> {
match stream {
#[cfg(unix)]
service_binding::Stream::Unix(stream) => {
let stream = tokio::net::UnixStream::from_std(stream)?;
Ok(Box::pin(Client::new(stream)))
Ok(Box::new(Client::new(stream)))
}
service_binding::Stream::Tcp(stream) => {
let stream = tokio::net::TcpStream::from_std(stream)?;
Ok(Box::pin(Client::new(stream)))
Ok(Box::new(Client::new(stream)))
}
#[cfg(windows)]
service_binding::Stream::NamedPipe(pipe) => {
Expand All @@ -65,9 +65,9 @@ pub async fn connect(
Err(e) => Err(e)?,
}

tokio::time::sleep(std::time::Duration::from_millis(50)).await;
std::thread::sleep(std::time::Duration::from_millis(50));
};
Ok(Box::pin(Client::new(stream)))
Ok(Box::new(Client::new(stream)))
}
#[cfg(not(windows))]
service_binding::Stream::NamedPipe(_) => Err(ProtoError::IO(std::io::Error::other(
Expand Down

0 comments on commit 1b556c8

Please sign in to comment.