Skip to content

Commit

Permalink
Implement the rest of the bore client
Browse files Browse the repository at this point in the history
  • Loading branch information
ekzhang committed Apr 6, 2022
1 parent fe1c8ad commit 599926d
Show file tree
Hide file tree
Showing 6 changed files with 245 additions and 76 deletions.
119 changes: 88 additions & 31 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 2 additions & 2 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,11 @@ name = "bore"
path = "src/main.rs"

[dependencies]
anyhow = "1.0.56"
anyhow = { version = "1.0.56", features = ["backtrace"] }
clap = { version = "3.1.8", features = ["derive"] }
dashmap = "5.2.0"
rmp-serde = "1.0.0"
serde = { version = "1.0.136", features = ["derive"] }
serde_json = "1.0.79"
tokio = { version = "1.17.0", features = ["full"] }
tracing = "0.1.32"
tracing-subscriber = "0.3.10"
Expand Down
96 changes: 96 additions & 0 deletions src/client.rs
Original file line number Diff line number Diff line change
@@ -1 +1,97 @@
//! Client implementation for the `bore` service.
use std::sync::Arc;

use anyhow::{bail, Context, Result};
use tokio::{io::BufReader, net::TcpStream};
use tracing::{error, info, info_span, warn, Instrument};
use uuid::Uuid;

use crate::shared::{proxy, recv_json, send_json, ClientMessage, ServerMessage, CONTROL_PORT};

/// State structure for the client.
pub struct Client {
/// Control connection to the server.
conn: Option<BufReader<TcpStream>>,

/// Destination address of the server.
to: String,

/// Local port that is forwarded.
local_port: u16,

/// Port that is publicly available on the remote.
remote_port: u16,
}

impl Client {
/// Create a new client.
pub async fn new(local_port: u16, to: &str, port: u16) -> Result<Self> {
let stream = TcpStream::connect((to, CONTROL_PORT)).await?;
let mut stream = BufReader::new(stream);

send_json(&mut stream, ClientMessage::Hello(port)).await?;
let remote_port = match recv_json(&mut stream, &mut Vec::new()).await? {
Some(ServerMessage::Hello(remote_port)) => remote_port,
Some(ServerMessage::Error(message)) => bail!("server error: {message}"),
Some(_) => bail!("unexpected initial non-hello message"),
None => bail!("unexpected EOF"),
};
info!(remote_port, "connected to server");
info!("listening at {to}:{remote_port}");

Ok(Client {
conn: Some(stream),
to: to.to_string(),
local_port,
remote_port,
})
}

/// Returns the port publicly available on the remote.
pub fn remote_port(&self) -> u16 {
self.remote_port
}

/// Start the client, listening for new connections.
pub async fn listen(mut self) -> Result<()> {
let mut conn = self.conn.take().unwrap();
let this = Arc::new(self);
let mut buf = Vec::new();
loop {
let msg = recv_json(&mut conn, &mut buf).await?;
match msg {
Some(ServerMessage::Hello(_)) => warn!("unexpected hello"),
Some(ServerMessage::Heartbeat) => (),
Some(ServerMessage::Connection(id)) => {
let this = Arc::clone(&this);
tokio::spawn(
async move {
info!("new connection");
match this.handle_connection(id).await {
Ok(_) => info!("connection exited"),
Err(err) => warn!(%err, "connection exited with error"),
}
}
.instrument(info_span!("proxy", %id)),
);
}
Some(ServerMessage::Error(err)) => error!(%err, "server error"),
None => return Ok(()),
}
}
}

async fn handle_connection(&self, id: Uuid) -> Result<()> {
let local_conn = TcpStream::connect(("localhost", self.local_port))
.await
.context("failed TCP connection to local port")?;
let mut remote_conn = TcpStream::connect((&self.to[..], CONTROL_PORT))
.await
.context("failed TCP connection to remote port")?;

send_json(&mut remote_conn, ClientMessage::Accept(id)).await?;
proxy(local_conn, remote_conn).await?;
Ok(())
}
}
6 changes: 3 additions & 3 deletions src/main.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use anyhow::Result;
use bore_cli::server::Server;
use bore_cli::{client::Client, server::Server};
use clap::{Parser, Subcommand};

#[derive(Parser, Debug)]
Expand Down Expand Up @@ -45,8 +45,8 @@ async fn main() -> Result<()> {
to,
port,
} => {
let _ = (local_port, to, port);
todo!()
let client = Client::new(local_port, &to, port).await?;
client.listen().await?;
}
Command::Server { min_port } => {
Server::new(min_port).listen().await?;
Expand Down
Loading

0 comments on commit 599926d

Please sign in to comment.