131 lines
4.5 KiB
Rust
131 lines
4.5 KiB
Rust
use std::{net::SocketAddr, path::Path};
|
|
|
|
use bytes::Bytes;
|
|
use http_body_util::{BodyExt, Full, combinators::BoxBody};
|
|
use hyper::{
|
|
Method, Request, Response, StatusCode,
|
|
client::conn::http1::Builder,
|
|
header::HOST,
|
|
server::conn::http1::{self},
|
|
service::service_fn,
|
|
};
|
|
use hyper_util::rt::TokioIo;
|
|
use rand::rngs::OsRng;
|
|
use tokio::net::TcpListener;
|
|
use tracing::{debug, trace, warn};
|
|
use tracing_subscriber::{EnvFilter, Registry, layer::SubscriberExt, util::SubscriberInitExt};
|
|
use tunnel_rs::ssh::Server;
|
|
|
|
fn full<T: Into<Bytes>>(chunk: T) -> BoxBody<Bytes, hyper::Error> {
|
|
Full::new(chunk.into())
|
|
.map_err(|never| match never {})
|
|
.boxed()
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
let env_filter = EnvFilter::try_from_default_env()
|
|
.or_else(|_| EnvFilter::try_new("info"))
|
|
.expect("Fallback should be valid");
|
|
|
|
let logger = tracing_subscriber::fmt::layer().compact();
|
|
Registry::default().with(logger).with(env_filter).init();
|
|
|
|
let key = if let Ok(path) = std::env::var("PRIVATE_KEY_FILE") {
|
|
russh::keys::PrivateKey::read_openssh_file(Path::new(&path)).unwrap()
|
|
} else {
|
|
russh::keys::PrivateKey::random(&mut OsRng, russh::keys::Algorithm::Ed25519).unwrap()
|
|
};
|
|
|
|
let mut ssh = Server::new();
|
|
let tunnels = ssh.tunnels();
|
|
tokio::spawn(async move { ssh.run(key, ("0.0.0.0", 2222)).await });
|
|
|
|
let service = service_fn(move |req: Request<_>| {
|
|
let tunnels = tunnels.clone();
|
|
async move {
|
|
if req.method() == Method::CONNECT {
|
|
let mut resp = Response::new(full("CONNECT not supported"));
|
|
*resp.status_mut() = StatusCode::BAD_REQUEST;
|
|
|
|
Ok::<_, hyper::Error>(resp)
|
|
} else {
|
|
trace!(?req);
|
|
let Some(authority) = req
|
|
.uri()
|
|
.authority()
|
|
.as_ref()
|
|
.map(|a| a.to_string())
|
|
.or_else(|| {
|
|
req.headers()
|
|
.get(HOST)
|
|
.map(|h| h.to_str().unwrap().to_owned())
|
|
})
|
|
else {
|
|
let mut resp = Response::new(full("Missing authority or host header"));
|
|
*resp.status_mut() = StatusCode::BAD_REQUEST;
|
|
|
|
return Ok::<_, hyper::Error>(resp);
|
|
};
|
|
|
|
debug!("Request for {authority:?}");
|
|
|
|
let Some(tunnel) = tunnels.read().await.get(&authority).cloned() else {
|
|
let mut resp = Response::new(full(format!("Unknown tunnel: {authority}")));
|
|
*resp.status_mut() = StatusCode::NOT_FOUND;
|
|
|
|
return Ok::<_, hyper::Error>(resp);
|
|
};
|
|
|
|
debug!("Opening channel");
|
|
let channel = match tunnel.open_tunnel().await {
|
|
Ok(channel) => channel,
|
|
Err(err) => {
|
|
warn!("Failed to tunnel: {err}");
|
|
let mut resp = Response::new(full("Failed to open tunnel"));
|
|
*resp.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
|
|
|
|
return Ok::<_, hyper::Error>(resp);
|
|
}
|
|
};
|
|
let io = TokioIo::new(channel.into_stream());
|
|
|
|
let (mut sender, conn) = Builder::new()
|
|
.preserve_header_case(true)
|
|
.title_case_headers(true)
|
|
.handshake(io)
|
|
.await?;
|
|
|
|
tokio::spawn(async move {
|
|
if let Err(err) = conn.await {
|
|
warn!("Connection failed: {err}");
|
|
}
|
|
});
|
|
|
|
let resp = sender.send_request(req).await.unwrap();
|
|
Ok(resp.map(|b| b.boxed()))
|
|
}
|
|
}
|
|
});
|
|
|
|
let addr = SocketAddr::from(([0, 0, 0, 0], 3000));
|
|
let listener = TcpListener::bind(addr).await.unwrap();
|
|
loop {
|
|
let (stream, _) = listener.accept().await.unwrap();
|
|
let io = TokioIo::new(stream);
|
|
|
|
let service = service.clone();
|
|
tokio::spawn(async move {
|
|
if let Err(err) = http1::Builder::new()
|
|
.preserve_header_case(true)
|
|
.title_case_headers(true)
|
|
.serve_connection(io, service)
|
|
.with_upgrades()
|
|
.await
|
|
{
|
|
warn!("Failed to serve connection: {err:?}");
|
|
}
|
|
});
|
|
}
|
|
}
|