Skip to content

Commit

Permalink
support multiple --restrict-http-upgrade-path-prefix
Browse files Browse the repository at this point in the history
  • Loading branch information
erebe committed Dec 3, 2023
1 parent 592fab9 commit 3eba7d8
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 9 deletions.
6 changes: 3 additions & 3 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -161,9 +161,9 @@ struct Server {
/// Server will only accept connection from if this specific path prefix is used during websocket upgrade.
/// Useful if you specify in the client a custom path prefix and you want the server to only allow this one.
/// The path prefix act as a secret to authenticate clients
/// Disabled by default. Accept all path prefix
/// Disabled by default. Accept all path prefix. Can be specified multiple time
#[arg(long, verbatim_doc_comment)]
restrict_http_upgrade_path_prefix: Option<String>,
restrict_http_upgrade_path_prefix: Option<Vec<String>>,

/// [Optional] Use custom certificate (.crt) instead of the default embedded self signed certificate.
#[arg(long, value_name = "FILE_PATH", verbatim_doc_comment)]
Expand Down Expand Up @@ -441,7 +441,7 @@ pub struct WsServerConfig {
pub socket_so_mark: Option<i32>,
pub bind: SocketAddr,
pub restrict_to: Option<Vec<String>>,
pub restrict_http_upgrade_path_prefix: Option<String>,
pub restrict_http_upgrade_path_prefix: Option<Vec<String>>,
pub websocket_ping_frequency: Option<Duration>,
pub timeout_connect: Duration,
pub websocket_mask_frame: bool,
Expand Down
9 changes: 6 additions & 3 deletions src/tunnel/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -202,12 +202,15 @@ async fn server_upgrade(
.unwrap_or_default());
}

if let Some(path_prefix) = &server_config.restrict_http_upgrade_path_prefix {
if let Some(paths_prefix) = &server_config.restrict_http_upgrade_path_prefix {
let path = req.uri().path();
let min_len = min(path.len(), 1);
let max_len = min(path.len(), path_prefix.len() + 1);
let mut max_len = 0;
if &path[0..min_len] != "/"
|| &path[min_len..max_len] != path_prefix.as_str()
|| !paths_prefix.iter().any(|p| {
max_len = min(path.len(), p.len() + 1);
p == &path[min_len..max_len]
})
|| !path[max_len..].starts_with('/')
{
warn!("Rejecting connection with bad path prefix in upgrade request: {}", req.uri());
Expand Down
14 changes: 11 additions & 3 deletions src/udp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -430,7 +430,9 @@ mod tests {
#[tokio::test]
async fn test_udp_server() {
let server_addr: SocketAddr = "[::1]:1234".parse().unwrap();
let server = run_server(server_addr, None, |_| Ok(()), |l| Ok(l.clone())).await.unwrap();
let server = run_server(server_addr, None, |_| Ok(()), |l| Ok(l.clone()))
.await
.unwrap();
pin_mut!(server);

// Should timeout
Expand Down Expand Up @@ -476,7 +478,11 @@ mod tests {
#[tokio::test]
async fn test_multiple_client() {
let server_addr: SocketAddr = "[::1]:1235".parse().unwrap();
let mut server = Box::pin(run_server(server_addr, None, |_| Ok(()), |l| Ok(l.clone())).await.unwrap());
let mut server = Box::pin(
run_server(server_addr, None, |_| Ok(()), |l| Ok(l.clone()))
.await
.unwrap(),
);

// Send some data to the server
let client = UdpSocket::bind("[::1]:0").await.unwrap();
Expand Down Expand Up @@ -542,7 +548,9 @@ mod tests {
async fn test_udp_should_timeout() {
let server_addr: SocketAddr = "[::1]:1237".parse().unwrap();
let socket_timeout = Duration::from_secs(1);
let server = run_server(server_addr, Some(socket_timeout), |_| Ok(()), |l| Ok(l.clone())).await.unwrap();
let server = run_server(server_addr, Some(socket_timeout), |_| Ok(()), |l| Ok(l.clone()))
.await
.unwrap();
pin_mut!(server);

// Send some data to the server
Expand Down

0 comments on commit 3eba7d8

Please sign in to comment.