From 48343b990feeaccfc261235c0dcbf947ba97c582 Mon Sep 17 00:00:00 2001 From: Eric Rodrigues Pires Date: Sat, 31 Jan 2026 21:36:31 -0300 Subject: [PATCH] Add `--pool-size` CLI flag Closes #77 --- CHANGELOG.md | 8 + Dockerfile | 2 +- book/src/advanced_options.md | 8 + book/src/configuration.md | 16 +- flake.lock | 10 +- src/config.rs | 14 +- src/entrypoint.rs | 10 + src/error.rs | 2 + src/http.rs | 59 ++++- src/lib.rs | 11 +- src/ssh/auth.rs | 7 +- src/ssh/connection_handler.rs | 87 ++++++- src/ssh/exec.rs | 30 +++ src/ssh/forwarding.rs | 18 +- src/ssh/mod.rs | 217 +++++++-------- tests/integration/alias_pool_limit.rs | 246 ++++++++++++++++++ tests/integration/config_invalid_options.rs | 33 ++- tests/integration/http_pool_limit.rs | 222 ++++++++++++++++ tests/integration/main.rs | 2 + .../integration/ssh_invalid_exec_commands.rs | 10 + 20 files changed, 860 insertions(+), 152 deletions(-) create mode 100644 tests/integration/alias_pool_limit.rs create mode 100644 tests/integration/http_pool_limit.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index cad6058..7da8b02 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,13 +4,21 @@ ### Added +- **BREAKING**: Add maximum pool size for proxy handlers. +- Add `--pool-size` CLI flag. - Add `--no-domain` CLI flag. +- Add `pool` option for remote forwarding connections. - Add keepalive mechanism for proxied HTTP connections. ### Fixed - Improve HTTP proxy performance. +### Changed + +- Lower `--directory-poll-interval` default value from 30 seconds to 15 seconds. +- Update Docker base image. + ## 0.8.7 (2026-01-15) ### Added diff --git a/Dockerfile b/Dockerfile index 4a4308a..5d4462d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,5 +1,5 @@ # Compile application with the official Rust image -FROM --platform=$BUILDPLATFORM rust:1.92.0-alpine3.22 AS builder +FROM --platform=$BUILDPLATFORM rust:1.93.0-alpine3.23 AS builder ENV PKGCONFIG_SYSROOTDIR=/ # Add build dependencies and targets RUN apk add --no-cache musl-dev libressl-dev perl build-base zig diff --git a/book/src/advanced_options.md b/book/src/advanced_options.md index 1d69ca8..9ea29c7 100644 --- a/book/src/advanced_options.md +++ b/book/src/advanced_options.md @@ -42,6 +42,14 @@ These options allow you to limit the IP ranges for incoming proxy/alias connecti ssh -p 2222 -R website.com:80:localhost:3000 sandhole.com.br ip-allowlist=10.0.0.0/8,20ff::/16 ip-blocklist=10.1.0.0/16 ``` +## `pool` + +This option allows you to reduce the maximum number of connections created for each of your handlers. It must not be larger than the server's `--pool-size` option. + +```bash +ssh -p 2222 -R my.tunnel:80:localhost:8080 sandhole.com.br pool=16 +``` + ## `sni-proxy` This option tells Sandhole that it should use your provided TLS backend. This guarantees that Sandhole cannot see unencrypted traffic. This option only works over HTTPS, so you may want to also set `force-https`. diff --git a/book/src/configuration.md b/book/src/configuration.md index 1c15161..d6b5403 100644 --- a/book/src/configuration.md +++ b/book/src/configuration.md @@ -14,23 +14,23 @@ Similarly, there is a `./deploy/admin_keys/` directory (set by `--admin-keys-dir Users with unrecognized SSH keys are still allowed to connect, in order to perform [local forwarding](./local_forwarding.mdrandom-subdomain-seed) to user-provided services. As such, these are the possible types of authentication: -| Authentication type | Connection method(s) | -| ------------------- | ----------------------- | -| None | Public key | -| User | Password¹ or public key | -| Admin | Public key | +| Authentication type | Connection method(s) | +| ------------------- | ---------------------------------- | +| None | Public key | +| User | Password1 or public key | +| Admin | Public key | -¹ Optional password authentication with a [login API](#alternative-authentication-with-password). +1 Optional password authentication with a [login API](#alternative-authentication-with-password). And these are each of their capabilities: | Authentication type | Local forwarding (proxy) | Remote forwading (reverse proxy) | Admin interface access | | ------------------- | ------------------------ | -------------------------------- | ---------------------- | | None | ✅ | ❌ | ❌ | -| User | ✅ | ✅² | ❌ | +| User | ✅ | ✅2 | ❌ | | Admin | ✅ | ✅ | ✅ | -² Remote forwarding by users is subject to restrictions, such as [service quotas](#service-quotas) and [rate limiting](#rate-limiting). +2 Remote forwarding by users is subject to restrictions, such as [service quotas](#service-quotas) and [rate limiting](#rate-limiting). ## Default ports diff --git a/flake.lock b/flake.lock index fd2fd86..e7edbaa 100644 --- a/flake.lock +++ b/flake.lock @@ -59,16 +59,14 @@ }, "rust-overlay": { "inputs": { - "nixpkgs": [ - "nixpkgs" - ] + "nixpkgs": ["nixpkgs"] }, "locked": { - "lastModified": 1768012928, - "narHash": "sha256-HFFVQaux1JoOjEvgBT0ASk1Je+jsipyO5c91FoLMed8=", + "lastModified": 1769828398, + "narHash": "sha256-zmnvRUm15QrlKH0V1BZoiT3U+Q+tr+P5Osi8qgtL9fY=", "owner": "oxalica", "repo": "rust-overlay", - "rev": "312b4371e72f644ffcff25b23615195e3b390643", + "rev": "a1d32c90c8a4ea43e9586b7e5894c179d5747425", "type": "github" }, "original": { diff --git a/src/config.rs b/src/config.rs index 0632b73..c534607 100644 --- a/src/config.rs +++ b/src/config.rs @@ -385,6 +385,13 @@ pub struct ApplicationConfig { )] pub buffer_size: u32, + /// Maximum pool size for simultaneous connections per proxied service. + /// The maximum is 1024. + /// + /// A higher value will lead to higher memory consumption, and may cause disruption on services. + #[arg(long, default_value_t = 128, value_name = "SIZE")] + pub pool_size: u16, + /// How long to wait between each keepalive message that is sent to an unresponsive SSH connection. #[arg(long, default_value = "15s", value_parser = validate_duration, value_name = "DURATION")] pub ssh_keepalive_interval: Duration, @@ -402,7 +409,7 @@ pub struct ApplicationConfig { /// A low value may consume too many resources on large file trees. #[arg( long, - default_value = "30s", + default_value = "15s", value_parser = validate_duration, value_name = "DURATION" )] @@ -551,9 +558,10 @@ mod application_config_tests { ip_allowlist: None, ip_blocklist: None, buffer_size: 32_768, + pool_size: 128, ssh_keepalive_interval: Duration::from_secs(15), ssh_keepalive_max: 3, - directory_poll_interval: Duration::from_secs(30), + directory_poll_interval: Duration::from_secs(15), idle_connection_timeout: Duration::from_secs(2), unproxied_connection_timeout: None, authentication_request_timeout: Duration::from_secs(5), @@ -610,6 +618,7 @@ mod application_config_tests { "--ip-allowlist=10.0.0.0/8", "--ip-blocklist=10.1.0.0/16,10.2.0.0/16", "--buffer-size=4KB", + "--pool-size=1024", "--ssh-keepalive-interval=10s", "--ssh-keepalive-max=2", "--directory-poll-interval=10s", @@ -671,6 +680,7 @@ mod application_config_tests { IpNet::from_str("10.2.0.0/16").unwrap() ]), buffer_size: 4_000, + pool_size: 1_024, ssh_keepalive_interval: Duration::from_secs(10), ssh_keepalive_max: 2, directory_poll_interval: Duration::from_secs(10), diff --git a/src/entrypoint.rs b/src/entrypoint.rs index d5bc070..4ae80c8 100644 --- a/src/entrypoint.rs +++ b/src/entrypoint.rs @@ -117,6 +117,12 @@ pub async fn entrypoint(config: ApplicationConfig) -> color_eyre::Result<()> { .into()); } } + let pool_size: usize = config.pool_size.into(); + if pool_size > 1024 { + return Err( + ServerError::InvalidConfig("Cannot set --pool-size greater than 1024".into()).into(), + ); + } let http_request_timeout = config.http_request_timeout; let tcp_connection_timeout = config.tcp_connection_timeout; let buffer_size = usize::try_from(config.buffer_size) @@ -617,6 +623,7 @@ pub async fn entrypoint(config: ApplicationConfig) -> color_eyre::Result<()> { }) // Always use aliasing channels instead of tunneling channels. .proxy_type(ProxyType::Aliasing) + .pool_size(pool_size) .buffer_size(buffer_size) .maybe_http_request_timeout(http_request_timeout) .maybe_websocket_timeout(tcp_connection_timeout) @@ -656,6 +663,7 @@ pub async fn entrypoint(config: ApplicationConfig) -> color_eyre::Result<()> { disable_tcp: config.disable_tcp, disable_aliasing: config.disable_aliasing, buffer_size, + pool_size, rate_limit: config .rate_limit_per_user .map(|rate| rate as f64) @@ -716,6 +724,7 @@ pub async fn entrypoint(config: ApplicationConfig) -> color_eyre::Result<()> { }) // Always use tunneling channels. .proxy_type(ProxyType::Tunneling) + .pool_size(pool_size) .buffer_size(buffer_size) .maybe_http_request_timeout(http_request_timeout) .maybe_websocket_timeout(tcp_connection_timeout) @@ -794,6 +803,7 @@ pub async fn entrypoint(config: ApplicationConfig) -> color_eyre::Result<()> { }) // Always use tunneling channels. .proxy_type(ProxyType::Tunneling) + .pool_size(pool_size) .buffer_size(buffer_size) .maybe_http_request_timeout(http_request_timeout) .maybe_websocket_timeout(tcp_connection_timeout) diff --git a/src/error.rs b/src/error.rs index 78727b9..93927d9 100644 --- a/src/error.rs +++ b/src/error.rs @@ -27,6 +27,8 @@ pub(crate) enum ServerError { TunnelingNotAllowed, #[error("Aliasing not allowed")] AliasingNotAllowed, + #[error("Pool limit reached")] + PoolLimitReached, #[error("SSH error: {0}")] Ssh(#[from] russh::Error), } diff --git a/src/http.rs b/src/http.rs index bf214fd..d3eb98b 100644 --- a/src/http.rs +++ b/src/http.rs @@ -132,6 +132,7 @@ struct HttpLog { elapsed_time: Duration, } +// Pretty-print an HTTP log with ANSI fn http_log(data: HttpLog, tx: Option, disable_http_logs: bool) { let HttpLog { ip, @@ -228,6 +229,8 @@ pub(crate) enum HttpError { HyperError(#[from] hyper::Error), #[error("Handler not found")] HandlerNotFound, + #[error("No handler available")] + NoHandlerAvailable, #[error("Channel request denied")] ChannelRequestDenied, #[error("Header to string error: {0}")] @@ -251,9 +254,13 @@ pub(crate) enum HttpError { } impl From for HttpError { - fn from(_: ServerError) -> Self { - // If getting the handler failed, return 404 (they may have an allowlist for fingerprints/IP networks) - HttpError::ChannelRequestDenied + fn from(value: ServerError) -> Self { + match value { + // If pool limit was reached, return 429 (no handler is available) + ServerError::PoolLimitReached => HttpError::NoHandlerAvailable, + // Otherwise, return 404 (they may have an allowlist for fingerprints/IP networks) + _ => HttpError::ChannelRequestDenied, + } } } @@ -270,7 +277,8 @@ impl IntoResponse for HttpError { | HttpError::InvalidUri(_) | HttpError::InvalidUriParts(_) | HttpError::MissingUpgradeHeader => StatusCode::BAD_REQUEST, - HttpError::RequestTimeout => StatusCode::REQUEST_TIMEOUT, + HttpError::NoHandlerAvailable => StatusCode::TOO_MANY_REQUESTS, + HttpError::RequestTimeout => StatusCode::GATEWAY_TIMEOUT, HttpError::HandlerNotFound | HttpError::ChannelRequestDenied => StatusCode::NOT_FOUND, HttpError::HyperError(_) => StatusCode::INTERNAL_SERVER_ERROR, } @@ -295,18 +303,21 @@ where ::Data: Send + Sync + 'static, ::Error: Error + Send + Sync + 'static, { + // Keep-alive pool for opened HTTP/1.1 connections. #[builder(default = DashMap::default())] keepalive_http11_pool_map: DashMap< KeepaliveAlias, KeepalivePool<(http1::SendRequest, ServerHandlerSender)>, RandomState, >, + // Keep-alive pool for opened HTTP/2 connections. #[builder(default = DashMap::default())] keepalive_http2_pool_map: DashMap< KeepaliveAlias, KeepalivePool<(http2::SendRequest, ServerHandlerSender)>, RandomState, >, + // Connection manager to get handlers from. conn_manager: M, // Tuple containing where to redirect requests from the main domain to. domain_redirect: Option>, @@ -314,6 +325,8 @@ where protocol: Protocol, // Configuration on which type of channel to retrieve from the handler. proxy_type: ProxyType, + // Pool size for connections reuse per handler. + pool_size: usize, // Buffer size for bidirectional copying. buffer_size: usize, // Optional duration until an outgoing request is canceled. @@ -642,7 +655,7 @@ where let elapsed_time = timer.elapsed(); http_log( http_log_builder - .status(StatusCode::REQUEST_TIMEOUT.as_u16()) + .status(StatusCode::GATEWAY_TIMEOUT.as_u16()) .elapsed_time(elapsed_time) .build(), Some(tx), @@ -695,7 +708,8 @@ where v.2.lock().expect("not poisoned").take(); }) .or_insert_with(|| { - let (sender, receiver) = async_channel::bounded(128); + let (sender, receiver) = + async_channel::bounded(proxy_data.pool_size); (sender, receiver, Arc::default()) }) .downgrade(); @@ -782,7 +796,7 @@ where let elapsed_time = timer.elapsed(); http_log( http_log_builder - .status(StatusCode::REQUEST_TIMEOUT.as_u16()) + .status(StatusCode::GATEWAY_TIMEOUT.as_u16()) .elapsed_time(elapsed_time) .build(), Some(tx), @@ -913,7 +927,7 @@ where let elapsed_time = timer.elapsed(); http_log( http_log_builder - .status(StatusCode::REQUEST_TIMEOUT.as_u16()) + .status(StatusCode::GATEWAY_TIMEOUT.as_u16()) .elapsed_time(elapsed_time) .build(), Some(tx), @@ -967,7 +981,8 @@ where v.2.lock().expect("not poisoned").take(); }) .or_insert_with(|| { - let (sender, receiver) = async_channel::bounded(128); + let (sender, receiver) = + async_channel::bounded(proxy_data.pool_size); (sender, receiver, Arc::default()) }) .downgrade(); @@ -1050,6 +1065,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Http { port: 80 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -1094,6 +1110,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Http { port: 80 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -1138,6 +1155,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Http { port: 80 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -1202,6 +1220,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::TlsRedirect { from: 80, to: 443 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -1269,6 +1288,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::TlsRedirect { from: 80, to: 8443 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -1337,6 +1357,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Http { port: 80 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -1404,6 +1425,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Http { port: 80 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -1497,6 +1519,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Https { port: 443 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .http_request_timeout(Duration::from_millis(500)) .disable_http_logs(false) @@ -1509,7 +1532,7 @@ mod proxy_handler_tests { "should log after timing out request" ); let response = response.expect("should return response after proxy"); - assert_eq!(response.status(), hyper::StatusCode::REQUEST_TIMEOUT); + assert_eq!(response.status(), hyper::StatusCode::GATEWAY_TIMEOUT); jh.abort(); } @@ -1586,6 +1609,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Https { port: 443 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .http_request_timeout(Duration::from_millis(500)) .disable_http_logs(false) @@ -1606,7 +1630,7 @@ mod proxy_handler_tests { match error { tokio_tungstenite::tungstenite::Error::Http(response) => { assert!( - response.status() == StatusCode::REQUEST_TIMEOUT, + response.status() == StatusCode::GATEWAY_TIMEOUT, "should've timed out Websocket request" ) } @@ -1697,6 +1721,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Https { port: 443 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .http_request_timeout(Duration::from_millis(500)) .disable_http_logs(false) @@ -1709,7 +1734,7 @@ mod proxy_handler_tests { "should log after timing out request" ); let response = response.expect("should return response after proxy"); - assert_eq!(response.status(), hyper::StatusCode::REQUEST_TIMEOUT); + assert_eq!(response.status(), hyper::StatusCode::GATEWAY_TIMEOUT); jh.abort(); } @@ -1796,6 +1821,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Https { port: 443 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -1896,6 +1922,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Https { port: 443 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -1995,6 +2022,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Https { port: 443 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -2093,6 +2121,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Http { port: 80 }) .proxy_type(ProxyType::Aliasing) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -2192,6 +2221,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Https { port: 443 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -2286,6 +2316,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Https { port: 443 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -2403,6 +2434,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Https { port: 443 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -2511,6 +2543,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Https { port: 443 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -2636,6 +2669,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Https { port: 443 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), @@ -2711,6 +2745,7 @@ mod proxy_handler_tests { })) .protocol(Protocol::Https { port: 443 }) .proxy_type(ProxyType::Tunneling) + .pool_size(128) .buffer_size(8_000) .disable_http_logs(false) .build(), diff --git a/src/lib.rs b/src/lib.rs index 2cf5475..9a2adcd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -13,9 +13,9 @@ use std::{ }; use ahash::RandomState; -use async_speed_limit::{Limiter, Resource, clock::StandardClock}; +use async_speed_limit::Limiter; use hyper::body::Incoming; -use russh::{ChannelStream, keys::ssh_key::Fingerprint, server::Msg}; +use russh::keys::ssh_key::Fingerprint; use tokio_util::sync::CancellationToken; use crate::{ @@ -26,7 +26,7 @@ use crate::{ http::ProxyData, quota::TokenHolderUser, reactor::{AliasReactor, HttpReactor, SniReactor, SshReactor, TcpReactor}, - ssh::connection_handler::SshTunnelHandler, + ssh::connection_handler::{SshChannel, SshTunnelHandler}, tcp::TcpHandler, tcp_alias::TcpAlias, }; @@ -82,8 +82,7 @@ type SessionMap = (Limiter, HashMap); // A generic table with data for the admin interface. type DataTable = Arc>>; // Helper type for HTTP proxy data types. -type HttpProxyData = - Arc, SshTunnelHandler, Resource, StandardClock>>>; +type HttpProxyData = Arc, SshTunnelHandler, SshChannel>>; // HTTP proxy data used by the tunneling connections. type TunnelingProxyData = HttpProxyData, HttpReactor>>; // HTTP proxy data used by the local forwarding aliasing connections. @@ -153,6 +152,8 @@ pub(crate) struct SandholeServer { pub(crate) disable_aliasing: bool, // Buffer size for bidirectional copying. pub(crate) buffer_size: usize, + // Pool size for SSH handlers. + pub(crate) pool_size: usize, // Rate limit per second for services of a single user. pub(crate) rate_limit: f64, // How long until a login API request is timed out. diff --git a/src/ssh/auth.rs b/src/ssh/auth.rs index fcbe37c..6d85c5b 100644 --- a/src/ssh/auth.rs +++ b/src/ssh/auth.rs @@ -3,7 +3,7 @@ use std::{ fmt::Display, sync::{ Arc, Mutex, RwLock, - atomic::{AtomicIsize, Ordering}, + atomic::{AtomicIsize, AtomicUsize, Ordering}, }, time::Duration, }; @@ -79,6 +79,8 @@ pub(crate) struct UserData { pub(crate) allow_fingerprint: Arc>>, // Extra data available for HTTP tunneling/aliasing connections. pub(crate) http_data: Arc>, + // Maximum amount of simultaneous connections for each handler. + pub(crate) max_pool_size: Arc, // Optional IP filtering for this connection's tunneling and aliasing channels. pub(crate) ip_filter: Arc>>, // What kind of restriction to impose on tunnels and aliases for this session. @@ -100,7 +102,7 @@ pub(crate) struct UserData { } impl UserData { - pub(crate) fn new(quota_key: TokenHolder, limiter: Limiter) -> Self { + pub(crate) fn new(quota_key: TokenHolder, limiter: Limiter, max_pool_size: usize) -> Self { Self { allow_fingerprint: Arc::new(RwLock::new(Box::new(|_| true))), http_data: Arc::new(RwLock::new(ConnectionHttpData { @@ -109,6 +111,7 @@ impl UserData { http2: false, host: None, })), + max_pool_size: Arc::new(AtomicUsize::new(max_pool_size)), ip_filter: Arc::new(RwLock::new(None)), session_restriction: UserSessionRestriction::None, quota_key, diff --git a/src/ssh/connection_handler.rs b/src/ssh/connection_handler.rs index c75fae9..75687c6 100644 --- a/src/ssh/connection_handler.rs +++ b/src/ssh/connection_handler.rs @@ -1,10 +1,15 @@ use std::{ net::{IpAddr, SocketAddr}, - sync::{Arc, RwLock}, + pin::pin, + sync::{ + Arc, RwLock, + atomic::{AtomicUsize, Ordering}, + }, }; use async_speed_limit::{Limiter, Resource, clock::StandardClock}; use russh::{ChannelStream, keys::ssh_key::Fingerprint, server::Msg}; +use tokio::io::{AsyncRead, AsyncWrite}; use crate::{ connection_handler::{ConnectionHandler, ConnectionHttpData}, @@ -13,6 +18,52 @@ use crate::{ ssh::{FingerprintFn, ServerHandlerSender}, }; +// Reference-counted wrapper of an SSH channel stream. +pub(crate) struct SshChannel { + inner: Resource, StandardClock>, + current_pool_size: Arc, +} + +impl Drop for SshChannel { + fn drop(&mut self) { + self.current_pool_size.fetch_sub(1, Ordering::Release); + } +} + +impl AsyncRead for SshChannel { + fn poll_read( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + buf: &mut tokio::io::ReadBuf<'_>, + ) -> std::task::Poll> { + pin!(&mut self.inner).poll_read(cx, buf) + } +} + +impl AsyncWrite for SshChannel { + fn poll_write( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + buf: &[u8], + ) -> std::task::Poll> { + pin!(&mut self.inner).poll_write(cx, buf) + } + + fn poll_flush( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + pin!(&mut self.inner).poll_flush(cx) + } + + fn poll_shutdown( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + pin!(&mut self.inner).poll_shutdown(cx) + } +} + // Struct for generating tunneling/aliasing channels from an underlying SSH connection, // via remote forwarding. It also includes a log channel to communicate messages // (such as HTTP logs) back to the SSH connection. @@ -22,6 +73,10 @@ pub(crate) struct SshTunnelHandler { pub(crate) allow_fingerprint: Arc>>, // Optional extra data available for HTTP tunneling/aliasing connections. pub(crate) http_data: Option>>, + // Maximum amount of simultaneous connections for this handler. + pub(crate) max_pool_size: Arc, + // Current active connections for this handler. + pub(crate) current_pool_size: Arc, // Optional IP filtering for this handler's tunneling and aliasing channels. pub(crate) ip_filter: Arc>>, // Handle to the SSH connection, in order to create remote forwarding channels. @@ -52,16 +107,12 @@ impl Drop for SshTunnelHandler { } } -impl ConnectionHandler, StandardClock>> for SshTunnelHandler { +impl ConnectionHandler for SshTunnelHandler { fn log_channel(&self) -> ServerHandlerSender { self.tx.clone() } - async fn tunneling_channel( - &self, - ip: IpAddr, - port: u16, - ) -> Result, StandardClock>, ServerError> { + async fn tunneling_channel(&self, ip: IpAddr, port: u16) -> Result { // Check if this IP is not blocked let tunneling_allowed = self .ip_filter @@ -70,6 +121,11 @@ impl ConnectionHandler, StandardClock>> for SshTunne .as_ref() .is_none_or(|filter| filter.is_allowed(ip)); if tunneling_allowed { + let max_pool_size = self.max_pool_size.load(Ordering::Acquire); + if self.current_pool_size.fetch_add(1, Ordering::AcqRel) >= max_pool_size { + self.current_pool_size.fetch_sub(1, Ordering::Release); + return Err(ServerError::PoolLimitReached); + } let channel = self .handle .channel_open_forwarded_tcpip( @@ -80,7 +136,10 @@ impl ConnectionHandler, StandardClock>> for SshTunne ) .await? .into_stream(); - Ok(self.limiter.clone().limit(channel)) + Ok(SshChannel { + inner: self.limiter.clone().limit(channel), + current_pool_size: Arc::clone(&self.current_pool_size), + }) } else { Err(ServerError::TunnelingNotAllowed) } @@ -102,8 +161,13 @@ impl ConnectionHandler, StandardClock>> for SshTunne ip: IpAddr, port: u16, fingerprint: Option<&'_ Fingerprint>, - ) -> Result, StandardClock>, ServerError> { + ) -> Result { if self.can_alias(ip, port, fingerprint) { + let max_pool_size = self.max_pool_size.load(Ordering::Acquire); + if self.current_pool_size.fetch_add(1, Ordering::AcqRel) >= max_pool_size { + self.current_pool_size.fetch_sub(1, Ordering::Release); + return Err(ServerError::PoolLimitReached); + } let channel = self .handle .channel_open_forwarded_tcpip( @@ -114,7 +178,10 @@ impl ConnectionHandler, StandardClock>> for SshTunne ) .await? .into_stream(); - Ok(self.limiter.clone().limit(channel)) + Ok(SshChannel { + inner: self.limiter.clone().limit(channel), + current_pool_size: Arc::clone(&self.current_pool_size), + }) } else { Err(ServerError::AliasingNotAllowed) } diff --git a/src/ssh/exec.rs b/src/ssh/exec.rs index 622b8e6..22475cf 100644 --- a/src/ssh/exec.rs +++ b/src/ssh/exec.rs @@ -25,6 +25,7 @@ pub(crate) enum ExecCommandFlag { IpAllowlist, IpBlocklist, Host, + Pool, } pub(crate) struct SshCommandContext<'a> { @@ -383,3 +384,32 @@ impl SshCommand for HostCommand { } } } + +pub(crate) struct PoolCommand(pub(crate) usize); + +impl SshCommand for PoolCommand { + fn flag(&self) -> ExecCommandFlag { + ExecCommandFlag::Pool + } + + async fn execute(&mut self, context: &mut SshCommandContext<'_>) -> color_eyre::Result<()> { + if context.commands.contains(self.flag()) { + return Err(eyre!("duplicated command")); + } + match context.auth_data { + AuthenticatedData::User { user_data } | AuthenticatedData::Admin { user_data, .. } => { + let prev_pool_size = user_data + .max_pool_size + .fetch_min(self.0, std::sync::atomic::Ordering::AcqRel); + if prev_pool_size < self.0 { + return Err(eyre!( + "pool size must be less than or equal to {prev_pool_size}" + )); + }; + context.commands.insert(self.flag()); + Ok(()) + } + _ => Err(eyre!("not authenticated as user")), + } + } +} diff --git a/src/ssh/forwarding.rs b/src/ssh/forwarding.rs index 8e10a9e..b71f85f 100644 --- a/src/ssh/forwarding.rs +++ b/src/ssh/forwarding.rs @@ -1,4 +1,8 @@ -use std::{borrow::Borrow, net::SocketAddr, sync::Arc}; +use std::{ + borrow::Borrow, + net::SocketAddr, + sync::{Arc, atomic::AtomicUsize}, +}; use chrono::Utc; use color_eyre::eyre::eyre; @@ -253,6 +257,8 @@ impl ForwardingHandlerStrategy for SshForwardingHandler { Arc::new(SshTunnelHandler { allow_fingerprint: Arc::clone(&context.user_data.allow_fingerprint), http_data: None, + max_pool_size: Arc::clone(&context.user_data.max_pool_size), + current_pool_size: Arc::new(AtomicUsize::new(0)), ip_filter: Arc::clone(&context.user_data.ip_filter), handle, tx: context.tx.clone(), @@ -526,6 +532,8 @@ impl ForwardingHandlerStrategy for HttpForwardingHandler { Arc::new(SshTunnelHandler { allow_fingerprint: Arc::clone(&context.user_data.allow_fingerprint), http_data: Some(Arc::clone(&context.user_data.http_data)), + max_pool_size: Arc::clone(&context.user_data.max_pool_size), + current_pool_size: Arc::new(AtomicUsize::new(0)), ip_filter: Arc::clone(&context.user_data.ip_filter), handle, tx: context.tx.clone(), @@ -596,6 +604,8 @@ impl ForwardingHandlerStrategy for HttpForwardingHandler { Arc::new(SshTunnelHandler { allow_fingerprint: Arc::clone(&context.user_data.allow_fingerprint), http_data: Some(Arc::clone(&context.user_data.http_data)), + max_pool_size: Arc::clone(&context.user_data.max_pool_size), + current_pool_size: Arc::new(AtomicUsize::new(0)), ip_filter: Arc::clone(&context.user_data.ip_filter), handle, tx: context.tx.clone(), @@ -705,6 +715,8 @@ impl ForwardingHandlerStrategy for HttpForwardingHandler { Arc::new(SshTunnelHandler { allow_fingerprint: Arc::clone(&context.user_data.allow_fingerprint), http_data: Some(Arc::clone(&context.user_data.http_data)), + max_pool_size: Arc::clone(&context.user_data.max_pool_size), + current_pool_size: Arc::new(AtomicUsize::new(0)), ip_filter: Arc::clone(&context.user_data.ip_filter), handle, tx: context.tx.clone(), @@ -1106,6 +1118,8 @@ impl ForwardingHandlerStrategy for AliasForwardingHandler { Arc::new(SshTunnelHandler { allow_fingerprint: Arc::clone(&context.user_data.allow_fingerprint), http_data: None, + max_pool_size: Arc::clone(&context.user_data.max_pool_size), + current_pool_size: Arc::new(AtomicUsize::new(0)), ip_filter: Arc::clone(&context.user_data.ip_filter), handle, tx: context.tx.clone(), @@ -1555,6 +1569,8 @@ impl ForwardingHandlerStrategy for TcpForwardingHandler { Arc::new(SshTunnelHandler { allow_fingerprint: Arc::clone(&context.user_data.allow_fingerprint), http_data: None, + max_pool_size: Arc::clone(&context.user_data.max_pool_size), + current_pool_size: Arc::new(AtomicUsize::new(0)), ip_filter: Arc::clone(&context.user_data.ip_filter), handle, tx: context.tx.clone(), diff --git a/src/ssh/mod.rs b/src/ssh/mod.rs index e9cc0ac..6f1fa0d 100644 --- a/src/ssh/mod.rs +++ b/src/ssh/mod.rs @@ -24,8 +24,8 @@ use crate::{ auth::{AdminData, AuthenticatedData, ProxyAutoCancellation, UserData}, exec::{ AdminCommand, AllowedFingerprintsCommand, ExecCommandFlag, ForceHttpsCommand, - HostCommand, Http2Command, IpAllowlistCommand, IpBlocklistCommand, SniProxyCommand, - SshCommand, SshCommandContext, TcpAliasCommand, + HostCommand, Http2Command, IpAllowlistCommand, IpBlocklistCommand, PoolCommand, + SniProxyCommand, SshCommand, SshCommandContext, TcpAliasCommand, }, forwarding::{Forwarder, LocalForwardingContext, RemoteForwardingContext}, }, @@ -241,6 +241,7 @@ impl Handler for ServerHandler { user_data: Box::new(UserData::new( TokenHolder::User(UserIdentification::Username(user.into())), limiter, + self.server.pool_size, )), }; #[cfg(not(coverage_nightly))] @@ -321,6 +322,7 @@ impl Handler for ServerHandler { user_data: Box::new(UserData::new( TokenHolder::User(UserIdentification::PublicKey(fingerprint)), limiter, + self.server.pool_size, )), }; } @@ -331,6 +333,7 @@ impl Handler for ServerHandler { user_data: Box::new(UserData::new( TokenHolder::Admin(UserIdentification::PublicKey(fingerprint)), limiter, + self.server.pool_size, )), admin_data: Box::new(AdminData::new()), }; @@ -416,7 +419,7 @@ impl Handler for ServerHandler { // - `admin` command creates an admin interface if the user is an admin "admin" => { let mut command = AdminCommand; - match command + if let Err(error) = command .execute(&mut SshCommandContext { server: &self.server, auth_data: &mut self.auth_data, @@ -426,17 +429,14 @@ impl Handler for ServerHandler { }) .await { - Ok(_) => {} - Err(error) => { - #[cfg(not(coverage_nightly))] - tracing::debug!(peer = %self.peer, %error, "'admin' command failed."); - let _ = self - .tx - .send(format!("'admin' command failed: {error}\r\n").into()); - success = false; - self.cancellation_token.cancel(); - break; - } + #[cfg(not(coverage_nightly))] + tracing::debug!(peer = %self.peer, %error, "'admin' command failed."); + let _ = self + .tx + .send(format!("'admin' command failed: {error}\r\n").into()); + success = false; + self.cancellation_token.cancel(); + break; } } // - `allowed-fingerprints` sets this connection as alias-only, @@ -449,7 +449,7 @@ impl Handler for ServerHandler { .map(|key| key.parse::()) .collect(); let mut command = AllowedFingerprintsCommand(set); - match command + if let Err(error) = command .execute(&mut SshCommandContext { server: &self.server, auth_data: &mut self.auth_data, @@ -459,24 +459,20 @@ impl Handler for ServerHandler { }) .await { - Ok(_) => {} - Err(error) => { - #[cfg(not(coverage_nightly))] - tracing::debug!(peer = %self.peer, %error, "'allowed-fingerprints' command failed."); - let _ = self.tx.send( - format!("'allowed-fingerprints' command failed: {error}\r\n") - .into(), - ); - success = false; - self.cancellation_token.cancel(); - break; - } + #[cfg(not(coverage_nightly))] + tracing::debug!(peer = %self.peer, %error, "'allowed-fingerprints' command failed."); + let _ = self.tx.send( + format!("'allowed-fingerprints' command failed: {error}\r\n").into(), + ); + success = false; + self.cancellation_token.cancel(); + break; } } // - `tcp-alias` sets this connection as alias-only. "tcp-alias" => { let mut command = TcpAliasCommand; - match command + if let Err(error) = command .execute(&mut SshCommandContext { server: &self.server, auth_data: &mut self.auth_data, @@ -486,23 +482,20 @@ impl Handler for ServerHandler { }) .await { - Ok(_) => {} - Err(error) => { - #[cfg(not(coverage_nightly))] - tracing::debug!(peer = %self.peer, %error, "'tcp-alias' command failed."); - let _ = self - .tx - .send(format!("'tcp-alias' command failed: {error}\r\n").into()); - success = false; - self.cancellation_token.cancel(); - break; - } + #[cfg(not(coverage_nightly))] + tracing::debug!(peer = %self.peer, %error, "'tcp-alias' command failed."); + let _ = self + .tx + .send(format!("'tcp-alias' command failed: {error}\r\n").into()); + success = false; + self.cancellation_token.cancel(); + break; } } // - `force-https` causes tunneled HTTP requests to be redirected to HTTPS. "force-https" => { let mut command = ForceHttpsCommand; - match command + if let Err(error) = command .execute(&mut SshCommandContext { server: &self.server, auth_data: &mut self.auth_data, @@ -512,23 +505,20 @@ impl Handler for ServerHandler { }) .await { - Ok(_) => {} - Err(error) => { - #[cfg(not(coverage_nightly))] - tracing::debug!(peer = %self.peer, %error, "'force-https' command failed."); - let _ = self - .tx - .send(format!("'force-https' command failed: {error}\r\n").into()); - success = false; - self.cancellation_token.cancel(); - break; - } + #[cfg(not(coverage_nightly))] + tracing::debug!(peer = %self.peer, %error, "'force-https' command failed."); + let _ = self + .tx + .send(format!("'force-https' command failed: {error}\r\n").into()); + success = false; + self.cancellation_token.cancel(); + break; } } // - `http2` allows serving HTTP/2 to the HTTP endpoints. "http2" => { let mut command = Http2Command; - match command + if let Err(error) = command .execute(&mut SshCommandContext { server: &self.server, auth_data: &mut self.auth_data, @@ -538,23 +528,20 @@ impl Handler for ServerHandler { }) .await { - Ok(_) => {} - Err(error) => { - #[cfg(not(coverage_nightly))] - tracing::debug!(peer = %self.peer, %error, "'http2' command failed."); - let _ = self - .tx - .send(format!("'http2' command failed: {error}\r\n").into()); - success = false; - self.cancellation_token.cancel(); - break; - } + #[cfg(not(coverage_nightly))] + tracing::debug!(peer = %self.peer, %error, "'http2' command failed."); + let _ = self + .tx + .send(format!("'http2' command failed: {error}\r\n").into()); + success = false; + self.cancellation_token.cancel(); + break; } } // - `sni-proxy` allows the user to handle the certificates themselves for HTTPS traffic. "sni-proxy" => { let mut command = SniProxyCommand; - match command + if let Err(error) = command .execute(&mut SshCommandContext { server: &self.server, auth_data: &mut self.auth_data, @@ -564,17 +551,14 @@ impl Handler for ServerHandler { }) .await { - Ok(_) => {} - Err(error) => { - #[cfg(not(coverage_nightly))] - tracing::debug!(peer = %self.peer, %error, "'sni-proxy' command failed."); - let _ = self - .tx - .send(format!("'sni-proxy' command failed: {error}\r\n").into()); - success = false; - self.cancellation_token.cancel(); - break; - } + #[cfg(not(coverage_nightly))] + tracing::debug!(peer = %self.peer, %error, "'sni-proxy' command failed."); + let _ = self + .tx + .send(format!("'sni-proxy' command failed: {error}\r\n").into()); + success = false; + self.cancellation_token.cancel(); + break; } } // - `ip-allowlist` requires tunneling/aliasing connections to come from @@ -587,7 +571,7 @@ impl Handler for ServerHandler { .map(|network| network.parse::()) .collect(); let mut command = IpAllowlistCommand(list); - match command + if let Err(error) = command .execute(&mut SshCommandContext { server: &self.server, auth_data: &mut self.auth_data, @@ -597,17 +581,14 @@ impl Handler for ServerHandler { }) .await { - Ok(_) => {} - Err(error) => { - #[cfg(not(coverage_nightly))] - tracing::debug!(peer = %self.peer, %error, "'ip-allowlist' command failed."); - let _ = self - .tx - .send(format!("'ip-allowlist' command failed: {error}\r\n").into()); - success = false; - self.cancellation_token.cancel(); - break; - } + #[cfg(not(coverage_nightly))] + tracing::debug!(peer = %self.peer, %error, "'ip-allowlist' command failed."); + let _ = self + .tx + .send(format!("'ip-allowlist' command failed: {error}\r\n").into()); + success = false; + self.cancellation_token.cancel(); + break; } } // - `ip-blocklist` requires tunneling/aliasing connections to come from @@ -620,7 +601,7 @@ impl Handler for ServerHandler { .map(|network| network.parse::()) .collect(); let mut command = IpBlocklistCommand(list); - match command + if let Err(error) = command .execute(&mut SshCommandContext { server: &self.server, auth_data: &mut self.auth_data, @@ -630,24 +611,21 @@ impl Handler for ServerHandler { }) .await { - Ok(_) => {} - Err(error) => { - #[cfg(not(coverage_nightly))] - tracing::debug!(peer = %self.peer, %error, "'ip-blocklist' command failed."); - let _ = self - .tx - .send(format!("'ip-blocklist' command failed: {error}\r\n").into()); - success = false; - self.cancellation_token.cancel(); - break; - } + #[cfg(not(coverage_nightly))] + tracing::debug!(peer = %self.peer, %error, "'ip-blocklist' command failed."); + let _ = self + .tx + .send(format!("'ip-blocklist' command failed: {error}\r\n").into()); + success = false; + self.cancellation_token.cancel(); + break; } } // - `host` changes the Host header of HTTP requests command if command.starts_with("host=") => { let host = command.trim_start_matches("host=").to_string(); let mut command = HostCommand(host); - match command + if let Err(error) = command .execute(&mut SshCommandContext { server: &self.server, auth_data: &mut self.auth_data, @@ -657,17 +635,50 @@ impl Handler for ServerHandler { }) .await { - Ok(_) => {} + #[cfg(not(coverage_nightly))] + tracing::debug!(peer = %self.peer, %error, "'host' command failed."); + let _ = self + .tx + .send(format!("'host' command failed: {error}\r\n").into()); + success = false; + self.cancellation_token.cancel(); + break; + } + } + // - `pool` decreases the pool size of the connection's handlers + command if command.starts_with("pool=") => { + let pool_size: usize = match command.trim_start_matches("pool=").parse() { + Ok(pool_size) => pool_size, Err(error) => { #[cfg(not(coverage_nightly))] - tracing::debug!(peer = %self.peer, %error, "'host' command failed."); + tracing::debug!(peer = %self.peer, %error, "'pool' command failed."); let _ = self .tx - .send(format!("'host' command failed: {error}\r\n").into()); + .send(format!("Invalid parameter for 'pool': {error}\r\n").into()); success = false; self.cancellation_token.cancel(); break; } + }; + let mut command = PoolCommand(pool_size); + if let Err(error) = command + .execute(&mut SshCommandContext { + server: &self.server, + auth_data: &mut self.auth_data, + peer: &self.peer, + commands: &mut self.commands, + tx: &self.tx, + }) + .await + { + #[cfg(not(coverage_nightly))] + tracing::debug!(peer = %self.peer, %error, "'pool' command failed."); + let _ = self + .tx + .send(format!("'pool' command failed: {error}\r\n").into()); + success = false; + self.cancellation_token.cancel(); + break; } } // - Unknown command diff --git a/tests/integration/alias_pool_limit.rs b/tests/integration/alias_pool_limit.rs new file mode 100644 index 0000000..6cb18f3 --- /dev/null +++ b/tests/integration/alias_pool_limit.rs @@ -0,0 +1,246 @@ +use std::{sync::Arc, time::Duration}; + +use clap::Parser; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha20Rng; +use russh::{ + Channel, + client::{Msg, Session}, + keys::ssh_key::private::Ed25519Keypair, +}; +use russh::{ + ChannelId, + keys::{key::PrivateKeyWithHashAlg, load_secret_key}, +}; +use sandhole::{ApplicationConfig, entrypoint}; +use tokio::{ + net::TcpStream, + sync::oneshot, + time::{sleep, timeout}, +}; + +use crate::common::SandholeHandle; + +/// This test ensures that no more connections than the specified pool limit +/// are able to connect at the same time. +#[test_log::test(tokio::test(flavor = "multi_thread"))] +async fn alias_pool_limit() { + // 1. Initialize Sandhole + let config = ApplicationConfig::parse_from([ + "sandhole", + "--domain=foobar.tld", + "--user-keys-directory", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/user_keys"), + "--admin-keys-directory", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/admin_keys"), + "--certificates-directory", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/certificates"), + "--private-key-file", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/server_keys/ssh"), + "--acme-cache-directory", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/acme_cache"), + "--disable-directory-creation", + "--listen-address=127.0.0.1", + "--ssh-port=18022", + "--http-port=18080", + "--https-port=18443", + "--acme-use-staging", + "--bind-hostnames=all", + "--idle-connection-timeout=1s", + "--authentication-request-timeout=5s", + "--http-request-timeout=10s", + "--pool-size=10", + ]); + let _sandhole_handle = SandholeHandle(tokio::spawn(async move { entrypoint(config).await })); + if timeout(Duration::from_secs(5), async { + while TcpStream::connect("127.0.0.1:18022").await.is_err() { + sleep(Duration::from_millis(100)).await; + } + }) + .await + .is_err() + { + panic!("Timeout waiting for Sandhole to start.") + }; + + // 2. Start SSH client that will be proxied + let key = load_secret_key( + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/private_keys/key1"), + None, + ) + .expect("Missing file key1"); + let (tx, rx) = oneshot::channel(); + let ssh_client = SshClient(Some(tx)); + let mut session = russh::client::connect(Default::default(), "127.0.0.1:18022", ssh_client) + .await + .expect("Failed to connect to SSH server"); + assert!( + session + .authenticate_publickey( + "user", + PrivateKeyWithHashAlg::new( + Arc::new(key), + session.best_supported_rsa_hash().await.unwrap().flatten() + ) + ) + .await + .expect("SSH authentication failed") + .success(), + "authentication didn't succeed" + ); + session + .tcpip_forward("some.alias", 12345) + .await + .expect("tcpip_forward failed"); + let channel = session + .channel_open_session() + .await + .expect("channel_open_session_failed"); + channel.exec(true, "pool=2").await.expect("exec failed"); + let Ok(channel_id) = timeout(Duration::from_secs(2), async { rx.await.unwrap() }).await else { + panic!("Timeout waiting for server to reply."); + }; + assert_eq!(channel_id, channel.id()); + + // 3. Start long-running requests that fill the pool + let mut jhs = Vec::new(); + for _ in 0..2 { + let key = russh::keys::PrivateKey::from(Ed25519Keypair::from_seed( + &ChaCha20Rng::from_os_rng().random(), + )); + let ssh_client = SshAliasClient; + let mut client_session = + russh::client::connect(Default::default(), "127.0.0.1:18022", ssh_client) + .await + .expect("Failed to connect to SSH server"); + assert!( + client_session + .authenticate_publickey( + "user", + PrivateKeyWithHashAlg::new( + Arc::new(key), + client_session + .best_supported_rsa_hash() + .await + .unwrap() + .flatten() + ) + ) + .await + .expect("SSH authentication failed") + .success(), + "authentication didn't succeed" + ); + let mut channel = client_session + .channel_open_direct_tcpip("some.alias", 12345, "::1", 23456) + .await + .expect("Local forwarding failed"); + let jh = tokio::spawn(async move { + while let Some(msg) = channel.wait().await { + if let russh::ChannelMsg::Data { data } = msg { + assert_eq!(&data[..], &b"Ping"[..]); + break; + } + } + drop(client_session); + }); + jhs.push(jh); + } + + // 3. Start request that gets rate-limited from pool exhaustion + tokio::time::sleep(Duration::from_millis(500)).await; + let key = russh::keys::PrivateKey::from(Ed25519Keypair::from_seed( + &ChaCha20Rng::from_os_rng().random(), + )); + let ssh_client = SshAliasClient; + let mut client_session = + russh::client::connect(Default::default(), "127.0.0.1:18022", ssh_client) + .await + .expect("Failed to connect to SSH server"); + assert!( + client_session + .authenticate_publickey( + "user", + PrivateKeyWithHashAlg::new( + Arc::new(key), + client_session + .best_supported_rsa_hash() + .await + .unwrap() + .flatten() + ) + ) + .await + .expect("SSH authentication failed") + .success(), + "authentication didn't succeed" + ); + assert!( + client_session + .channel_open_direct_tcpip("some.alias", 12345, "::1", 23456) + .await + .is_err() + ); + timeout(Duration::from_secs(5), async move { + for jh in jhs { + jh.await.unwrap(); + } + }) + .await + .expect("timeout waiting for join handles to finish"); +} + +struct SshClient(Option>); + +impl russh::client::Handler for SshClient { + type Error = color_eyre::eyre::Error; + + async fn check_server_key( + &mut self, + _key: &russh::keys::PublicKey, + ) -> Result { + Ok(true) + } + + async fn server_channel_open_forwarded_tcpip( + &mut self, + channel: Channel, + _connected_address: &str, + _connected_port: u32, + _originator_address: &str, + _originator_port: u32, + _session: &mut Session, + ) -> Result<(), Self::Error> { + tokio::spawn(async move { + tokio::time::sleep(Duration::from_secs(3)).await; + channel.data(&b"Ping"[..]).await.unwrap(); + channel.eof().await.unwrap(); + channel.close().await.unwrap(); + }); + Ok(()) + } + + async fn channel_success( + &mut self, + channel: ChannelId, + _session: &mut Session, + ) -> Result<(), Self::Error> { + if let Some(tx) = self.0.take() { + tx.send(channel).unwrap(); + }; + Ok(()) + } +} + +struct SshAliasClient; + +impl russh::client::Handler for SshAliasClient { + type Error = color_eyre::eyre::Error; + + async fn check_server_key( + &mut self, + _key: &russh::keys::PublicKey, + ) -> Result { + Ok(true) + } +} diff --git a/tests/integration/config_invalid_options.rs b/tests/integration/config_invalid_options.rs index 79d296a..7ffd686 100644 --- a/tests/integration/config_invalid_options.rs +++ b/tests/integration/config_invalid_options.rs @@ -75,7 +75,36 @@ async fn config_invalid_options() { panic!("Timeout waiting for Sandhole to start.") }; - // 3a. Fail to initialize ACME ALPN resolver if HTTPS port is not 443 + // 3. Fail to initialize with --pool-size greater than 1024 + let config = ApplicationConfig::parse_from([ + "sandhole", + "--no-domain", + "--user-keys-directory", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/user_keys"), + "--admin-keys-directory", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/admin_keys"), + "--certificates-directory", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/certificates"), + "--private-key-file", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/server_keys/ssh"), + "--acme-cache-directory", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/acme_cache"), + "--listen-address=127.0.0.1", + "--ssh-port=18022", + "--http-port=18080", + "--https-port=18443", + "--pool-size=1025", + ]); + if timeout(Duration::from_secs(5), async { + assert!(entrypoint(config).await.is_err()); + }) + .await + .is_err() + { + panic!("Timeout waiting for Sandhole to start.") + }; + + // 4a. Fail to initialize ACME ALPN resolver if HTTPS port is not 443 let config = ApplicationConfig::parse_from([ "sandhole", "--domain=foobar.tld", @@ -108,7 +137,7 @@ async fn config_invalid_options() { { panic!("Timeout waiting for Sandhole to start.") }; - // 3b. Fail to connect with fake TLS-ALPN-01 challenge verifier + // 4b. Fail to connect with fake TLS-ALPN-01 challenge verifier let mut tls_config = ClientConfig::builder() .dangerous() .with_custom_certificate_verifier(Arc::new(SkipServerVerification)) diff --git a/tests/integration/http_pool_limit.rs b/tests/integration/http_pool_limit.rs new file mode 100644 index 0000000..f84176e --- /dev/null +++ b/tests/integration/http_pool_limit.rs @@ -0,0 +1,222 @@ +use std::{sync::Arc, time::Duration}; + +use axum::{Router, extract::Request, routing::get}; +use clap::Parser; +use http::{StatusCode, header::HOST}; +use http_body_util::BodyExt; +use hyper::{body::Incoming, service::service_fn}; +use hyper_util::{ + rt::{TokioExecutor, TokioIo}, + server::conn::auto::Builder, +}; +use russh::keys::{key::PrivateKeyWithHashAlg, load_secret_key}; +use russh::{ + Channel, + client::{Msg, Session}, +}; +use sandhole::{ApplicationConfig, entrypoint}; +use tokio::{ + net::TcpStream, + time::{sleep, timeout}, +}; +use tower::Service; + +use crate::common::SandholeHandle; + +/// This test ensures that no more connections than the specified pool limit +/// are able to connect at the same time via HTTP. +#[test_log::test(tokio::test(flavor = "multi_thread"))] +async fn http_pool_limit() { + // 1. Initialize Sandhole + let config = ApplicationConfig::parse_from([ + "sandhole", + "--domain=foobar.tld", + "--user-keys-directory", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/user_keys"), + "--admin-keys-directory", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/admin_keys"), + "--certificates-directory", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/certificates"), + "--private-key-file", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/server_keys/ssh"), + "--acme-cache-directory", + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/acme_cache"), + "--disable-directory-creation", + "--listen-address=127.0.0.1", + "--ssh-port=18022", + "--http-port=18080", + "--https-port=18443", + "--acme-use-staging", + "--bind-hostnames=all", + "--idle-connection-timeout=1s", + "--authentication-request-timeout=5s", + "--http-request-timeout=10s", + "--pool-size=2", + ]); + let _sandhole_handle = SandholeHandle(tokio::spawn(async move { entrypoint(config).await })); + if timeout(Duration::from_secs(5), async { + while TcpStream::connect("127.0.0.1:18022").await.is_err() { + sleep(Duration::from_millis(100)).await; + } + }) + .await + .is_err() + { + panic!("Timeout waiting for Sandhole to start.") + }; + + // 2. Start SSH client that will be proxied + let key = load_secret_key( + concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/private_keys/key1"), + None, + ) + .expect("Missing file key1"); + let ssh_client = SshClient; + let mut session = russh::client::connect(Default::default(), "127.0.0.1:18022", ssh_client) + .await + .expect("Failed to connect to SSH server"); + assert!( + session + .authenticate_publickey( + "user", + PrivateKeyWithHashAlg::new( + Arc::new(key), + session.best_supported_rsa_hash().await.unwrap().flatten() + ) + ) + .await + .expect("SSH authentication failed") + .success(), + "authentication didn't succeed" + ); + session + .tcpip_forward("test.foobar.tld", 80) + .await + .expect("tcpip_forward failed"); + + // 3. Start long-running requests that fill the pool + let mut jhs = Vec::new(); + for _ in 0..2 { + let tcp_stream = TcpStream::connect("127.0.0.1:18080") + .await + .expect("TCP connection failed"); + let (mut sender, conn) = hyper::client::conn::http1::handshake(TokioIo::new(tcp_stream)) + .await + .expect("HTTP handshake failed"); + let request = Request::builder() + .method("GET") + .uri("/") + .header(HOST, "test.foobar.tld") + .body(http_body_util::Empty::::new()) + .unwrap(); + let jh = tokio::spawn(async move { + let jh = tokio::spawn(async move { + if let Err(error) = conn.await { + eprintln!("Connection failed: {error:?}"); + } + }); + let Ok(response) = timeout(Duration::from_secs(10), async move { + sender + .send_request(request) + .await + .expect("Error sending HTTP request") + }) + .await + else { + panic!("Timeout waiting for request to finish."); + }; + assert_eq!(response.status(), StatusCode::OK); + let response_body = String::from_utf8( + response + .into_body() + .collect() + .await + .expect("Error collecting response") + .to_bytes() + .into(), + ) + .expect("Invalid response body"); + assert_eq!(response_body, "Processed"); + jh.abort(); + }); + jhs.push(jh); + } + + // 3. Start request that gets rate-limited from pool exhaustion + tokio::time::sleep(Duration::from_millis(500)).await; + let tcp_stream = TcpStream::connect("127.0.0.1:18080") + .await + .expect("TCP connection failed"); + let (mut sender, conn) = hyper::client::conn::http1::handshake(TokioIo::new(tcp_stream)) + .await + .expect("HTTP handshake failed"); + let jh = tokio::spawn(async move { + if let Err(error) = conn.await { + eprintln!("Connection failed: {error:?}"); + } + }); + let request = Request::builder() + .method("GET") + .uri("/") + .header(HOST, "test.foobar.tld") + .body(http_body_util::Empty::::new()) + .unwrap(); + let Ok(response) = timeout(Duration::from_secs(5), async move { + sender + .send_request(request) + .await + .expect("Error sending HTTP request") + }) + .await + else { + panic!("Timeout waiting for request to finish."); + }; + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + jh.abort(); + timeout(Duration::from_secs(5), async move { + for jh in jhs { + jh.await.unwrap(); + } + }) + .await + .expect("timeout waiting for join handles to finish"); +} + +struct SshClient; + +impl russh::client::Handler for SshClient { + type Error = color_eyre::eyre::Error; + + async fn check_server_key( + &mut self, + _key: &russh::keys::PublicKey, + ) -> Result { + Ok(true) + } + + async fn server_channel_open_forwarded_tcpip( + &mut self, + channel: Channel, + _connected_address: &str, + _connected_port: u32, + _originator_address: &str, + _originator_port: u32, + _session: &mut Session, + ) -> Result<(), Self::Error> { + let router = Router::new().route( + "/", + get(async || { + tokio::time::sleep(Duration::from_secs(3)).await; + "Processed" + }), + ); + let service = service_fn(move |req: Request| router.clone().call(req)); + tokio::spawn(async move { + Builder::new(TokioExecutor::new()) + .serve_connection_with_upgrades(TokioIo::new(channel.into_stream()), service) + .await + .expect("Invalid request"); + }); + Ok(()) + } +} diff --git a/tests/integration/main.rs b/tests/integration/main.rs index ea7e595..33da869 100644 --- a/tests/integration/main.rs +++ b/tests/integration/main.rs @@ -12,6 +12,7 @@ mod alias_aliasing_tunnel; mod alias_cannot_be_localhost; mod alias_http_aliases; mod alias_local_forward_existing_http; +mod alias_pool_limit; mod alias_rate_limit_download; mod alias_rate_limit_upload; mod alias_reject_special_ports; @@ -41,6 +42,7 @@ mod http_addressing_profanities_domain; mod http_addressing_profanities_no_root_domain; mod http_addressing_profanities_subdomain; mod http_force_https_by_user; +mod http_pool_limit; mod http_redirects; mod http_rewrite_host; mod http_simultaneous_requests; diff --git a/tests/integration/ssh_invalid_exec_commands.rs b/tests/integration/ssh_invalid_exec_commands.rs index d6f7dd0..e18cdd7 100644 --- a/tests/integration/ssh_invalid_exec_commands.rs +++ b/tests/integration/ssh_invalid_exec_commands.rs @@ -40,6 +40,7 @@ async fn ssh_invalid_exec_commands() { "--idle-connection-timeout=2s", "--authentication-request-timeout=5s", "--http-request-timeout=5s", + "--pool-size=64", ]); let _sandhole_handle = SandholeHandle(tokio::spawn(async move { entrypoint(config).await })); if timeout(Duration::from_secs(5), async { @@ -111,6 +112,15 @@ async fn ssh_invalid_exec_commands() { // `host` twice vec!["host=hello.world host=goodbye.world"], vec!["host=hello.world", "host=goodbye.world"], + // Invalid `pool` + vec!["pool=hello"], + // No `pool` + vec!["pool="], + // `pool` greater than the default + vec!["pool=65"], + // `pool` twice + vec!["pool=10 pool=10"], + vec!["pool=10", "pool=10"], // Unknown command vec!["unknown"], ] { -- 2.51.2