hex/wstunnel
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Cleanup

Browse source
This commit is contained in:
Σrebe - Romain GERARD 2023-10-30 08:13:38 +01:00
parent b478288848
commit 466cb425bc
No known key found for this signature in database
GPG key ID: 7A42B4B97E0332F4
11 changed files with 159 additions and 320 deletions
Download patch file Download diff file Expand all files Collapse all files

16
src/tunnel/client.rs
View file

@ -44,9 +44,7 @@ pub async fn connect(
) -> anyhow::Result<WebSocket<Upgraded>> {
let mut pooled_cnx = match client_cfg.cnx_pool().get().await {
Ok(tcp_stream) => tcp_stream,
Err(err) => Err(anyhow!(
"failed to get a connection to the server from the pool: {err:?}"
))?,
Err(err) => Err(anyhow!("failed to get a connection to the server from the pool: {err:?}"))?,
};
let mut req = Request::builder()
@ -80,12 +78,7 @@ pub async fn connect(
let transport = pooled_cnx.deref_mut().take().unwrap();
let (ws, _) = fastwebsockets::handshake::client(&SpawnExecutor, req, transport)
.await
.with_context(|| {
format!(
"failed to do websocket handshake with the server {:?}",
client_cfg.remote_addr
)
})?;
.with_context(|| format!("failed to do websocket handshake with the server {:?}", client_cfg.remote_addr))?;
Ok(ws)
}
@ -109,10 +102,7 @@ where
// Forward local tx to websocket tx
let ping_frequency = client_cfg.websocket_ping_frequency;
tokio::spawn(
super::io::propagate_read(local_rx, ws_tx, close_tx, ping_frequency)
.instrument(Span::current()),
);
tokio::spawn(super::io::propagate_read(local_rx, ws_tx, close_tx, ping_frequency).instrument(Span::current()));
// Forward websocket rx to local rx
let _ = super::io::propagate_write(local_tx, ws_rx, close_rx).await;

47
src/tunnel/io.rs
View file

@ -1,12 +1,12 @@
use fastwebsockets::{Frame, OpCode, Payload, WebSocketError, WebSocketRead, WebSocketWrite};
use futures_util::pin_mut;
use hyper::upgrade::Upgraded;
use std::pin::Pin;
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadHalf, WriteHalf};
use tokio::select;
use tokio::sync::oneshot;
use tokio::time::timeout;
use tracing::log::debug;
use tracing::{error, info, trace, warn};
@ -20,7 +20,14 @@ pub(super) async fn propagate_read(
info!("Closing local tx ==> websocket tx tunnel");
});
let mut buffer = vec![0u8; 8 * 1024];
static JUMBO_FRAME_SIZE: usize = 9 * 1024; // enough for a jumbo frame
let mut buffer = vec![0u8; JUMBO_FRAME_SIZE];
// We do our own pin_mut! to avoid shadowing timeout and be able to reset it, on next loop iteration
// We reuse the future to avoid creating a timer in the tight loop
let mut timeout_unpin = tokio::time::sleep(ping_frequency);
let mut timeout = unsafe { Pin::new_unchecked(&mut timeout_unpin) };
pin_mut!(local_rx);
loop {
let read_len = select! {
@ -30,9 +37,12 @@ pub(super) async fn propagate_read(
_ = close_tx.closed() => break,
_ = timeout(ping_frequency, futures_util::future::pending::<()>()) => {
_ = &mut timeout => {
debug!("sending ping to keep websocket connection alive");
ws_tx.write_frame(Frame::new(true, OpCode::Ping, None, Payload::BorrowedMut(&mut []))).await?;
timeout_unpin = tokio::time::sleep(ping_frequency);
timeout = unsafe { Pin::new_unchecked(&mut timeout_unpin) };
continue;
}
};
@ -41,32 +51,30 @@ pub(super) async fn propagate_read(
Ok(0) => break,
Ok(read_len) => read_len,
Err(err) => {
warn!(
"error while reading incoming bytes from local tx tunnel {}",
err
);
warn!("error while reading incoming bytes from local tx tunnel {}", err);
break;
}
};
trace!("read {} bytes", read_len);
match ws_tx
if let Err(err) = ws_tx
.write_frame(Frame::binary(Payload::BorrowedMut(&mut buffer[..read_len])))
.await
{
Ok(_) => {}
Err(err) => {
warn!("error while writing to websocket tx tunnel {}", err);
break;
}
warn!("error while writing to websocket tx tunnel {}", err);
break;
}
// If the buffer has been completely filled with previous read, Double it !
// For the buffer to not be a bottleneck when the TCP window scale
// For udp, the buffer will never grows.
if buffer.capacity() == read_len {
buffer.clear();
buffer.resize(buffer.capacity() * 2, 0);
}
}
// Send normal close
let _ = ws_tx.write_frame(Frame::close(1000, &[])).await;
Ok(())
@ -104,20 +112,15 @@ pub(super) async fn propagate_write(
trace!("receive ws frame {:?} {:?}", msg.opcode, msg.payload);
let ret = match msg.opcode {
OpCode::Continuation | OpCode::Text | OpCode::Binary => {
local_tx.write_all(msg.payload.as_ref()).await
}
OpCode::Continuation | OpCode::Text | OpCode::Binary => local_tx.write_all(msg.payload.as_ref()).await,
OpCode::Close => break,
OpCode::Ping => Ok(()),
OpCode::Pong => Ok(()),
};
match ret {
Ok(_) => {}
Err(err) => {
error!("error while writing bytes to local for rx tunnel {}", err);
break;
}
if let Err(err) = ret {
error!("error while writing bytes to local for rx tunnel {}", err);
break;
}
}

20
src/tunnel/mod.rs
View file

@ -42,12 +42,8 @@ impl JwtTunnelConfig {
}
static JWT_SECRET: &[u8; 15] = b"champignonfrais";
static JWT_KEY: Lazy<(Header, EncodingKey)> = Lazy::new(|| {
(
Header::new(Algorithm::HS256),
EncodingKey::from_secret(JWT_SECRET),
)
});
static JWT_KEY: Lazy<(Header, EncodingKey)> =
Lazy::new(|| (Header::new(Algorithm::HS256), EncodingKey::from_secret(JWT_SECRET)));
static JWT_DECODE: Lazy<(Validation, DecodingKey)> = Lazy::new(|| {
let mut validation = Validation::new(Algorithm::HS256);
@ -61,11 +57,7 @@ pub enum TransportStream {
}
impl AsyncRead for TransportStream {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
match self.get_mut() {
TransportStream::Plain(cnx) => Pin::new(cnx).poll_read(cx, buf),
TransportStream::Tls(cnx) => Pin::new(cnx).poll_read(cx, buf),
@ -74,11 +66,7 @@ impl AsyncRead for TransportStream {
}
impl AsyncWrite for TransportStream {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, Error>> {
fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, Error>> {
match self.get_mut() {
TransportStream::Plain(cnx) => Pin::new(cnx).poll_write(cx, buf),
TransportStream::Tls(cnx) => Pin::new(cnx).poll_write(cx, buf),

57
src/tunnel/server.rs
View file

@ -46,15 +46,8 @@ async fn from_query(
Span::current().record("remote", format!("{}:{}", jwt.claims.r, jwt.claims.rp));
if let Some(allowed_dests) = &server_config.restrict_to {
let requested_dest = format!("{}:{}", jwt.claims.r, jwt.claims.rp);
if allowed_dests
.iter()
.any(|dest| dest == &requested_dest)
.not()
{
warn!(
"Rejecting connection with not allowed destination: {}",
requested_dest
);
if allowed_dests.iter().any(|dest| dest == &requested_dest).not() {
warn!("Rejecting connection with not allowed destination: {}", requested_dest);
return Err(anyhow::anyhow!("Invalid upgrade request"));
}
}
@ -75,14 +68,9 @@ async fn from_query(
LocalProtocol::Tcp { .. } => {
let host = Host::parse(&jwt.claims.r)?;
let port = jwt.claims.rp;
let (rx, tx) = tcp::connect(
&host,
port,
&server_config.socket_so_mark,
Duration::from_secs(10),
)
.await?
.into_split();
let (rx, tx) = tcp::connect(&host, port, &server_config.socket_so_mark, Duration::from_secs(10))
.await?
.into_split();
Ok((jwt.claims.p, host, port, Box::pin(rx), Box::pin(tx)))
}
@ -100,10 +88,7 @@ async fn server_upgrade(
}
if !req.uri().path().ends_with("/events") {
warn!(
"Rejecting connection with bad upgrade request: {}",
req.uri()
);
warn!("Rejecting connection with bad upgrade request: {}", req.uri());
return Ok(http::Response::builder()
.status(StatusCode::BAD_REQUEST)
.body(Body::from("Invalid upgrade request"))
@ -118,10 +103,7 @@ async fn server_upgrade(
|| &path[min_len..max_len] != path_prefix.as_str()
|| !path[max_len..].starts_with('/')
{
warn!(
"Rejecting connection with bad path prefix in upgrade request: {}",
req.uri()
);
warn!("Rejecting connection with bad path prefix in upgrade request: {}", req.uri());
return Ok(http::Response::builder()
.status(StatusCode::BAD_REQUEST)
.body(Body::from("Invalid upgrade request"))
@ -133,11 +115,7 @@ async fn server_upgrade(
match from_query(&server_config, req.uri().query().unwrap_or_default()).await {
Ok(ret) => ret,
Err(err) => {
warn!(
"Rejecting connection with bad upgrade request: {} {}",
err,
req.uri()
);
warn!("Rejecting connection with bad upgrade request: {} {}", err, req.uri());
return Ok(http::Response::builder()
.status(StatusCode::BAD_REQUEST)
.body(Body::from(format!("Invalid upgrade request: {:?}", err)))
@ -149,11 +127,7 @@ async fn server_upgrade(
let (response, fut) = match fastwebsockets::upgrade::upgrade(&mut req) {
Ok(ret) => ret,
Err(err) => {
warn!(
"Rejecting connection with bad upgrade request: {} {}",
err,
req.uri()
);
warn!("Rejecting connection with bad upgrade request: {} {}", err, req.uri());
return Ok(http::Response::builder()
.status(StatusCode::BAD_REQUEST)
.body(Body::from(format!("Invalid upgrade request: {:?}", err)))
@ -171,14 +145,10 @@ async fn server_upgrade(
}
};
let (close_tx, close_rx) = oneshot::channel::<()>();
let ping_frequency = server_config
.websocket_ping_frequency
.unwrap_or(Duration::MAX);
let ping_frequency = server_config.websocket_ping_frequency.unwrap_or(Duration::MAX);
ws_tx.set_auto_apply_mask(server_config.websocket_mask_frame);
tokio::task::spawn(
super::io::propagate_write(local_tx, ws_rx, close_rx).instrument(Span::current()),
);
tokio::task::spawn(super::io::propagate_write(local_tx, ws_rx, close_rx).instrument(Span::current()));
let _ = super::io::propagate_read(local_rx, ws_tx, close_tx, ping_frequency).await;
}
@ -189,10 +159,7 @@ async fn server_upgrade(
}
pub async fn run_server(server_config: Arc<WsServerConfig>) -> anyhow::Result<()> {
info!(
"Starting wstunnel server listening on {}",
server_config.bind
);
info!("Starting wstunnel server listening on {}", server_config.bind);
let config = server_config.clone();
let upgrade_fn = move |req: Request<Body>| server_upgrade(config.clone(), req);