chore: add tests for command line parsing

2024-12-15 15:13:31 +01:00 · 2024-12-15 15:13:31 +01:00 · fab129b516
commit fab129b516
parent 314c619b8b
4 changed files with 474 additions and 1 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@ -471,6 +471,12 @@ dependencies = [
 "cc",
 ]

+[[package]]
+name = "collection_macros"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "50b180e6a75e306052a61658f832b4fc565a6e5a204da05f0fe7f50a31fb827a"
+
 [[package]]
 name = "colorchoice"
 version = "1.0.3"
@ -887,6 +893,12 @@ version = "0.3.31"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988"

+[[package]]
+name = "futures-timer"
+version = "3.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24"
+
 [[package]]
 name = "futures-util"
 version = "0.3.31"
@ -1945,6 +1957,15 @@ dependencies = [
 "syn",
 ]

+[[package]]
+name = "proc-macro-crate"
+version = "3.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8ecf48c7ca261d60b74ab1a7b20da18bede46776b2e55535cb958eb595c5fa7b"
+dependencies = [
+ "toml_edit",
+]
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.92"
@ -2074,6 +2095,12 @@ version = "0.8.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c"

+[[package]]
+name = "relative-path"
+version = "1.9.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2"
+
 [[package]]
 name = "resolv-conf"
 version = "0.7.0"
@ -2099,6 +2126,36 @@ dependencies = [
 "windows-sys 0.52.0",
 ]

+[[package]]
+name = "rstest"
+version = "0.23.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0a2c585be59b6b5dd66a9d2084aa1d8bd52fbdb806eafdeffb52791147862035"
+dependencies = [
+ "futures",
+ "futures-timer",
+ "rstest_macros",
+ "rustc_version",
+]
+
+[[package]]
+name = "rstest_macros"
+version = "0.23.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "825ea780781b15345a146be27eaefb05085e337e869bff01b4306a4fd4a9ad5a"
+dependencies = [
+ "cfg-if",
+ "glob",
+ "proc-macro-crate",
+ "proc-macro2",
+ "quote",
+ "regex",
+ "relative-path",
+ "rustc_version",
+ "syn",
+ "unicode-ident",
+]
+
 [[package]]
 name = "rustc-demangle"
 version = "0.1.24"
@ -2111,6 +2168,15 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"

+[[package]]
+name = "rustc_version"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92"
+dependencies = [
+ "semver",
+]
+
 [[package]]
 name = "rusticata-macros"
 version = "4.1.0"
@ -2259,6 +2325,15 @@ dependencies = [
 "winapi-util",
 ]

+[[package]]
+name = "scc"
+version = "2.2.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "66b202022bb57c049555430e11fc22fea12909276a80a4c3d368da36ac1d88ed"
+dependencies = [
+ "sdd",
+]
+
 [[package]]
 name = "schannel"
 version = "0.1.27"
@ -2284,6 +2359,12 @@ dependencies = [
 "untrusted 0.9.0",
 ]

+[[package]]
+name = "sdd"
+version = "3.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "49c1eeaf4b6a87c7479688c6d52b9f1153cedd3c489300564f932b065c6eab95"
+
 [[package]]
 name = "security-framework"
 version = "2.11.1"
@ -2320,6 +2401,12 @@ dependencies = [
 "libc",
 ]

+[[package]]
+name = "semver"
+version = "1.0.24"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3cb6eb87a131f756572d7fb904f6e7b68633f09cca868c5df1c4b8d1a694bbba"
+
 [[package]]
 name = "serde"
 version = "1.0.216"
@ -2428,6 +2515,31 @@ dependencies = [
 "unsafe-libyaml",
 ]

+[[package]]
+name = "serial_test"
+version = "3.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1b258109f244e1d6891bf1053a55d63a5cd4f8f4c30cf9a1280989f80e7a1fa9"
+dependencies = [
+ "futures",
+ "log",
+ "once_cell",
+ "parking_lot",
+ "scc",
+ "serial_test_derive",
+]
+
+[[package]]
+name = "serial_test_derive"
+version = "3.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5d69265a08751de7844521fd15003ae0a888e035773ba05695c5c759a6f89eef"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "sha1"
 version = "0.10.6"
@ -2584,6 +2696,39 @@ dependencies = [
 "syn",
 ]

+[[package]]
+name = "test-case"
+version = "3.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eb2550dd13afcd286853192af8601920d959b14c401fcece38071d53bf0768a8"
+dependencies = [
+ "test-case-macros",
+]
+
+[[package]]
+name = "test-case-core"
+version = "3.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adcb7fd841cd518e279be3d5a3eb0636409487998a4aff22f3de87b81e88384f"
+dependencies = [
+ "cfg-if",
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "test-case-macros"
+version = "3.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5c89e72a01ed4c579669add59014b9a524d609c0c88c6a585ce37485879f6ffb"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+ "test-case-core",
+]
+
 [[package]]
 name = "testcontainers"
 version = "0.23.1"
@ -2799,6 +2944,23 @@ dependencies = [
 "tokio",
 ]

+[[package]]
+name = "toml_datetime"
+version = "0.6.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0dd7358ecb8fc2f8d014bf86f6f638ce72ba252a2c3a2572f2a795f1d23efb41"
+
+[[package]]
+name = "toml_edit"
+version = "0.22.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4ae48d6208a266e853d946088ed816055e556cc6028c5e8e2b84d9fa5dd7c7f5"
+dependencies = [
+ "indexmap 2.7.0",
+ "toml_datetime",
+ "winnow",
+]
+
 [[package]]
 name = "tower-service"
 version = "0.3.3"
@ -3253,6 +3415,15 @@ version = "0.52.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"

+[[package]]
+name = "winnow"
+version = "0.6.20"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "36c1fec1a2bb5866f07c25f68c26e565c4c200aebb96d7e55710c19d3e8ac49b"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "winreg"
 version = "0.50.0"
@ -3287,6 +3458,7 @@ dependencies = [
 "bb8",
 "bytes",
 "clap",
+ "collection_macros",
 "crossterm",
 "fast-socks5",
 "fastwebsockets",
@ -3306,13 +3478,16 @@ dependencies = [
 "ppp",
 "rcgen",
 "regex",
+ "rstest",
 "rustls-native-certs 0.8.1",
 "rustls-pemfile 2.2.0",
 "scopeguard",
 "serde",
 "serde_regex",
 "serde_yaml",
+ "serial_test",
 "socket2",
+ "test-case",
 "testcontainers",
 "tokio",
 "tokio-fd",
--- a/Cargo.toml
+++ b/Cargo.toml
@ -71,6 +71,10 @@ rcgen = { version = "0.13.1", default-features = false, features = ["ring"] }

 [dev-dependencies]
 testcontainers = "0.23.1"
+test-case = "3.3.1"
+collection_macros = "0.2.0"
+rstest = "0.23.0"
+serial_test = "3.2.0"

 [profile.release]
 lto = "fat"
--- a/src/main.rs
+++ b/src/main.rs
@ -1,6 +1,8 @@
 mod embedded_certificate;
 mod protocols;
 mod restrictions;
+#[cfg(test)]
+mod test_integrations;
 mod tunnel;

 use crate::protocols::dns::DnsResolver;
@ -379,7 +381,7 @@ struct Server {
    http_proxy_password: Option<String>,
 }

-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq)]
 pub struct LocalToRemote {
    local_protocol: LocalProtocol,
    local: SocketAddr,
@ -1200,3 +1202,62 @@ fn mk_http_proxy(

    Ok(Some(proxy))
 }
+
+#[cfg(test)]
+mod test {
+    use crate::tunnel::LocalProtocol;
+    use crate::{parse_local_bind, parse_tunnel_arg, parse_tunnel_dest, LocalToRemote};
+    use collection_macros::btreemap;
+    use std::collections::BTreeMap;
+    use std::io;
+    use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
+    use test_case::test_case;
+    use url::Host;
+
+    #[test_case("localhost:443" => (Host::Domain("localhost".to_string()), 443, BTreeMap::new()) ; "with domain")]
+    #[test_case("127.0.0.1:443" => (Host::Ipv4(Ipv4Addr::new(127, 0, 0, 1)), 443, BTreeMap::new()) ; "with IPv4")]
+    #[test_case("[::1]:8080" => (Host::Ipv6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), 8080, BTreeMap::new()) ; "with IpV6")]
+    #[test_case("a:1?timeout_sec=30&b=5" => (Host::Domain("a".to_string()), 1, btreemap! { "b".to_string() => "5".to_string(), "timeout_sec".to_string() => "30".to_string() }) ; "with options")]
+    fn test_parse_tunnel_dest(input: &str) -> (Host<String>, u16, BTreeMap<String, String>) {
+        parse_tunnel_dest(input).unwrap()
+    }
+
+    const LOCALHOST_IP4: SocketAddrV4 = SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 443);
+    const LOCALHOST_IP6: SocketAddrV6 = SocketAddrV6::new(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1), 443, 0, 0);
+
+    #[test_case("domain.com:443" => matches Err(_) ; "with domain")]
+    #[test_case("127.0.0.1" => matches Err(_) ; "with no port")]
+    #[test_case("127.0.0.1:444444443" => matches Err(_) ; "with too long port")]
+    #[test_case("127.0.0.1:443" => matches Ok((SocketAddr::V4(LOCALHOST_IP4), _)) ; "with ipv4")]
+    #[test_case("[::1]:443" => matches Ok((SocketAddr::V6(LOCALHOST_IP6), _)) ; "with ipv6")]
+    fn test_parse_local_bind(input: &str) -> Result<(SocketAddr, &str), io::Error> {
+        parse_local_bind(input)
+    }
+
+    #[test_case("domain.com:443" => panics ""; "with no protocol")]
+    #[test_case("sdsf://443:domain.com:443" => panics ""; "with invalid protocol")]
+    #[test_case("tcp://443:domain.com:4443" =>
+        LocalToRemote {
+            local_protocol: LocalProtocol::Tcp { proxy_protocol: false },
+            local: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 443)),
+            remote: (Host::Domain("domain.com".to_string()), 4443),
+        }
+    ; "with no local bind")]
+    #[test_case("udp://[::1]:443:toto.com:4443?timeout_sec=30" =>
+        LocalToRemote {
+            local_protocol: LocalProtocol::Udp { timeout: Some(std::time::Duration::from_secs(30)) },
+            local: SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1), 443, 0, 0)),
+            remote: (Host::Domain("toto.com".to_string()), 4443),
+        }
+    ; "with fully defined tunnel")]
+    #[test_case("udp://[::1]:443:[::1]:4443?timeout_sec=30" =>
+        LocalToRemote {
+            local_protocol: LocalProtocol::Udp { timeout: Some(std::time::Duration::from_secs(30)) },
+            local: SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1), 443, 0, 0)),
+            remote: (Host::Ipv6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), 4443),
+        }
+    ; "with full ipv6 tunnel")]
+    fn test_parse_tunnel_arg(input: &str) -> LocalToRemote {
+        parse_tunnel_arg(input).unwrap()
+    }
+}
--- a/src/test_integrations.rs
+++ b/src/test_integrations.rs
@ -0,0 +1,233 @@
+use crate::protocols;
+use crate::protocols::dns::DnsResolver;
+use crate::restrictions::types;
+use crate::restrictions::types::{AllowConfig, MatchConfig, RestrictionConfig, RestrictionsRules};
+use crate::tunnel::client::{WsClient, WsClientConfig};
+use crate::tunnel::listeners::{TcpTunnelListener, UdpTunnelListener};
+use crate::tunnel::server::{WsServer, WsServerConfig};
+use crate::tunnel::transport::{TransportAddr, TransportScheme};
+use bytes::BytesMut;
+use futures_util::StreamExt;
+use hyper::http::HeaderValue;
+use ipnet::{IpNet, Ipv4Net, Ipv6Net};
+use regex::Regex;
+use rstest::{fixture, rstest};
+use scopeguard::defer;
+use serial_test::serial;
+use std::collections::HashMap;
+use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
+use std::time::Duration;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::pin;
+use url::Host;
+
+#[fixture]
+fn dns_resolver() -> DnsResolver {
+    DnsResolver::new_from_urls(&[], None, None, true).expect("Cannot create DNS resolver")
+}
+
+#[fixture]
+fn server_no_tls(dns_resolver: DnsResolver) -> WsServer {
+    let server_config = WsServerConfig {
+        socket_so_mark: None,
+        bind: "127.0.0.1:8080".parse().unwrap(),
+        websocket_ping_frequency: Some(Duration::from_secs(10)),
+        timeout_connect: Duration::from_secs(10),
+        websocket_mask_frame: false,
+        tls: None,
+        dns_resolver,
+        restriction_config: None,
+        http_proxy: None,
+    };
+    WsServer::new(server_config)
+}
+
+#[fixture]
+async fn client_ws(dns_resolver: DnsResolver) -> WsClient {
+    let client_config = WsClientConfig {
+        remote_addr: TransportAddr::new(TransportScheme::Ws, Host::Ipv4("127.0.0.1".parse().unwrap()), 8080, None)
+            .unwrap(),
+        socket_so_mark: None,
+        http_upgrade_path_prefix: "wstunnel".to_string(),
+        http_upgrade_credentials: None,
+        http_headers: HashMap::new(),
+        http_headers_file: None,
+        http_header_host: HeaderValue::from_static("127.0.0.1:8080"),
+        timeout_connect: Duration::from_secs(10),
+        websocket_ping_frequency: Some(Duration::from_secs(10)),
+        websocket_mask_frame: false,
+        dns_resolver,
+        http_proxy: None,
+    };
+
+    WsClient::new(client_config, 1, Duration::from_secs(1)).await.unwrap()
+}
+
+#[fixture]
+fn no_restrictions() -> RestrictionsRules {
+    pub fn default_host() -> Regex {
+        Regex::new("^.*$").unwrap()
+    }
+
+    pub fn default_cidr() -> Vec<IpNet> {
+        vec![IpNet::V4(Ipv4Net::default()), IpNet::V6(Ipv6Net::default())]
+    }
+
+    let tunnels = types::AllowConfig::Tunnel(types::AllowTunnelConfig {
+        protocol: vec![],
+        port: vec![],
+        host: default_host(),
+        cidr: default_cidr(),
+    });
+    let reverse_tunnel = AllowConfig::ReverseTunnel(types::AllowReverseTunnelConfig {
+        protocol: vec![],
+        port: vec![],
+        port_mapping: Default::default(),
+        cidr: default_cidr(),
+    });
+
+    RestrictionsRules {
+        restrictions: vec![RestrictionConfig {
+            name: "".to_string(),
+            r#match: vec![MatchConfig::Any],
+            allow: vec![tunnels, reverse_tunnel],
+        }],
+    }
+}
+
+const TUNNEL_LISTEN: (SocketAddr, Host) = (
+    SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 9998)),
+    Host::Ipv4(Ipv4Addr::new(127, 0, 0, 1)),
+);
+const ENDPOINT_LISTEN: (SocketAddr, Host) = (
+    SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 9999)),
+    Host::Ipv4(Ipv4Addr::new(127, 0, 0, 1)),
+);
+
+#[rstest]
+#[timeout(Duration::from_secs(10))]
+#[tokio::test]
+#[serial]
+async fn test_tcp_tunnel(
+    #[future] client_ws: WsClient,
+    server_no_tls: WsServer,
+    no_restrictions: RestrictionsRules,
+    dns_resolver: DnsResolver,
+) {
+    let server_h = tokio::spawn(server_no_tls.serve(no_restrictions));
+    defer! { server_h.abort(); };
+
+    let client_ws = client_ws.await;
+
+    let server = TcpTunnelListener::new(TUNNEL_LISTEN.0, (ENDPOINT_LISTEN.1, ENDPOINT_LISTEN.0.port()), false)
+        .await
+        .unwrap();
+    tokio::spawn(async move {
+        client_ws.run_tunnel(server).await.unwrap();
+    });
+
+    let mut tcp_listener = protocols::tcp::run_server(ENDPOINT_LISTEN.0, false).await.unwrap();
+    let mut client = protocols::tcp::connect(
+        &TUNNEL_LISTEN.1,
+        TUNNEL_LISTEN.0.port(),
+        None,
+        Duration::from_secs(10),
+        &dns_resolver,
+    )
+    .await
+    .unwrap();
+
+    client.write_all(b"Hello").await.unwrap();
+    let mut dd = tcp_listener.next().await.unwrap().unwrap();
+    let mut buf = BytesMut::new();
+    dd.read_buf(&mut buf).await.unwrap();
+    assert_eq!(&buf[..5], b"Hello");
+    buf.clear();
+
+    dd.write_all(b"world!").await.unwrap();
+    client.read_buf(&mut buf).await.unwrap();
+    assert_eq!(&buf[..6], b"world!");
+}
+
+#[rstest]
+#[timeout(Duration::from_secs(10))]
+#[tokio::test]
+#[serial]
+async fn test_udp_tunnel(
+    #[future] client_ws: WsClient,
+    server_no_tls: WsServer,
+    no_restrictions: RestrictionsRules,
+    dns_resolver: DnsResolver,
+) {
+    let server_h = tokio::spawn(server_no_tls.serve(no_restrictions));
+    defer! { server_h.abort(); };
+
+    let client_ws = client_ws.await;
+
+    let server = UdpTunnelListener::new(TUNNEL_LISTEN.0, (ENDPOINT_LISTEN.1, ENDPOINT_LISTEN.0.port()), None)
+        .await
+        .unwrap();
+    tokio::spawn(async move {
+        client_ws.run_tunnel(server).await.unwrap();
+    });
+
+    let udp_listener = protocols::udp::run_server(ENDPOINT_LISTEN.0, None, |_| Ok(()), |s| Ok(s.clone()))
+        .await
+        .unwrap();
+    let mut client = protocols::udp::connect(
+        &TUNNEL_LISTEN.1,
+        TUNNEL_LISTEN.0.port(),
+        Duration::from_secs(10),
+        None,
+        &dns_resolver,
+    )
+    .await
+    .unwrap();
+
+    client.write_all(b"Hello").await.unwrap();
+    pin!(udp_listener);
+    let dd = udp_listener.next().await.unwrap().unwrap();
+    pin!(dd);
+    let mut buf = BytesMut::new();
+    dd.read_buf(&mut buf).await.unwrap();
+    assert_eq!(&buf[..5], b"Hello");
+    buf.clear();
+
+    dd.writer().write_all(b"world!").await.unwrap();
+    client.read_buf(&mut buf).await.unwrap();
+    assert_eq!(&buf[..6], b"world!");
+}
+
+//#[rstest]
+//#[timeout(Duration::from_secs(10))]
+//#[tokio::test]
+//async fn test_socks5_tunnel(
+//    #[future] client_ws: WsClient,
+//    server_no_tls: WsServer,
+//    no_restrictions: RestrictionsRules,
+//    dns_resolver: DnsResolver,
+//) {
+//    let server_h = tokio::spawn(server_no_tls.serve(no_restrictions));
+//    defer! { server_h.abort(); };
+//
+//    let client_ws = client_ws.await;
+//
+//    let server = Socks5TunnelListener::new(TUNNEL_LISTEN.0, None, None).await.unwrap();
+//    tokio::spawn(async move { client_ws.run_tunnel(server).await.unwrap(); });
+//
+//    let socks5_listener = protocols::socks5::run_server(ENDPOINT_LISTEN.0, None, None).await.unwrap();
+//    let mut client = protocols::tcp::connect(&TUNNEL_LISTEN.1, TUNNEL_LISTEN.0.port(), None, Duration::from_secs(10), &dns_resolver).await.unwrap();
+//
+//    client.write_all(b"Hello").await.unwrap();
+//    pin!(socks5_listener);
+//    let (dd, _) = socks5_listener.next().await.unwrap().unwrap();
+//    let (mut read, mut write) = dd.into_split();
+//    let mut buf = BytesMut::new();
+//    read.read_buf(&mut buf).await.unwrap();
+//    assert_eq!(&buf[..5], b"Hello");
+//    buf.clear();
+//
+//    write.write_all(b"world!").await.unwrap();
+//    client.read_buf(&mut buf).await.unwrap();
+//    assert_eq!(&buf[..6], b"world!");
+//}