Mirrors/bore
1
0
Fork 0
mirror of https://github.com/ekzhang/bore.git synced 2025-12-16 03:47:50 +01:00
Code Issues Packages Projects Releases Wiki Activity

Improve stability by exiting immediately on common errors (#2)

Browse Source 
* Kill connections immediately on missing or close

* Add timeout to initial protocol messages

* Add low-level tracing for JSON messages

* Add timeout to initial TCP connections
This commit is contained in:
Eric Zhang
2022-04-08 15:55:54 -04:00
committed by GitHub
parent c1efefeddf
commit 2d0dcf9889
4 changed files with 51 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
src/client.rs
View File

@@ -3,12 +3,15 @@
use std::sync::Arc;
use anyhow::{bail, Context, Result};
use tokio::{io::BufReader, net::TcpStream};
use tokio::{io::BufReader, net::TcpStream, time::timeout};
use tracing::{error, info, info_span, warn, Instrument};
use uuid::Uuid;
use crate::auth::Authenticator;
use crate::shared::{proxy, recv_json, send_json, ClientMessage, ServerMessage, CONTROL_PORT};
use crate::shared::{
proxy, recv_json, recv_json_timeout, send_json, ClientMessage, ServerMessage, CONTROL_PORT,
NETWORK_TIMEOUT,
};
/// State structure for the client.
pub struct Client {
@@ -31,10 +34,7 @@ pub struct Client {
impl Client {
/// Create a new client.
pub async fn new(local_port: u16, to: &str, port: u16, secret: Option<&str>) -> Result<Self> {
let stream = TcpStream::connect((to, CONTROL_PORT))
.await
.with_context(|| format!("could not connect to {to}:{CONTROL_PORT}"))?;
let mut stream = BufReader::new(stream);
let mut stream = BufReader::new(connect_with_timeout(to, CONTROL_PORT).await?);
let auth = secret.map(Authenticator::new);
if let Some(auth) = &auth {
@@ -42,7 +42,7 @@ impl Client {
}
send_json(&mut stream, ClientMessage::Hello(port)).await?;
let remote_port = match recv_json(&mut stream, &mut Vec::new()).await? {
let remote_port = match recv_json_timeout(&mut stream).await? {
Some(ServerMessage::Hello(remote_port)) => remote_port,
Some(ServerMessage::Error(message)) => bail!("server error: {message}"),
Some(ServerMessage::Challenge(_)) => {
@@ -99,21 +99,23 @@ impl Client {
}
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 = BufReader::new(
TcpStream::connect((&self.to[..], CONTROL_PORT))
.await
.context("failed TCP connection to remote port")?,
);
let mut remote_conn =
BufReader::new(connect_with_timeout(&self.to[..], CONTROL_PORT).await?);
if let Some(auth) = &self.auth {
auth.client_handshake(&mut remote_conn).await?;
}
send_json(&mut remote_conn, ClientMessage::Accept(id)).await?;
let local_conn = connect_with_timeout("localhost", self.local_port).await?;
proxy(local_conn, remote_conn).await?;
Ok(())
}
}
async fn connect_with_timeout(to: &str, port: u16) -> Result<TcpStream> {
match timeout(NETWORK_TIMEOUT, TcpStream::connect((to, port))).await {
Ok(res) => res,
Err(err) => Err(err.into()),
}
.with_context(|| format!("could not connect to {to}:{port}"))
}