tproxy: The ending s in the url must be on the port part

Author: kvinwang

tproxy/src/proxy.rs

@@ -50,7 +50,7 @@ fn is_subdomain(sni: &str, base_domain: &str) -> bool {
 
 struct DstInfo {
     app_id: String,
-    port: Option<u16>,
+    port: u16,
     is_tls: bool,
 }
 
@@ -62,23 +62,38 @@ fn parse_destination(sni: &str, dotted_base_domain: &str) -> Result<DstInfo> {
     if subdomain.contains('.') {
         bail!("only one level of subdomain is supported");
     }
+    let mut parts = subdomain.split('-');
+    let app_id = parts.next().context("no app id found")?.to_owned();
+    if app_id.is_empty() {
+        bail!("app id is empty");
+    }
+    let last_part = parts.next();
     let is_tls;
-    let subdomain = match subdomain.strip_suffix("s") {
+    let port;
+    match last_part {
         None => {
             is_tls = false;
-            subdomain
+            port = None;
         }
-        Some(subdomain) => {
-            is_tls = true;
-            subdomain
+        Some(last_part) => {
+            let port_str = match last_part.strip_suffix('s') {
+                None => {
+                    is_tls = false;
+                    last_part
+                }
+                Some(last_part) => {
+                    is_tls = true;
+                    last_part
+                }
+            };
+            port = if port_str.is_empty() {
+                None
+            } else {
+                Some(port_str.parse::<u16>().context("invalid port")?)
+            };
         }
     };
-    let mut parts = subdomain.split('-');
-    let app_id = parts.next().context("no app id found")?.to_owned();
-    let port = parts
-        .next()
-        .map(|p| p.parse().context("invalid port"))
-        .transpose()?;
+    let port = port.unwrap_or(if is_tls { 443 } else { 80 });
     if parts.next().is_some() {
         bail!("invalid sni format");
     }
@@ -106,15 +121,9 @@ async fn handle_connection(
     if is_subdomain(&sni, dotted_base_domain) {
         let dst = parse_destination(&sni, dotted_base_domain)?;
         if dst.is_tls {
-            tls_passthough::proxy_to_app(
-                state,
-                inbound,
-                buffer,
-                &dst.app_id,
-                dst.port.unwrap_or(443),
-            )
-            .await
-            .with_context(|| format!("error on connection {sni}"))
+            tls_passthough::proxy_to_app(state, inbound, buffer, &dst.app_id, dst.port)
+                .await
+                .with_context(|| format!("error on connection {sni}"))
         } else {
             tls_terminate_proxy
                 .proxy(inbound, buffer, &dst.app_id, dst.port)
@@ -200,7 +209,70 @@ pub fn start(config: ProxyConfig, app_state: Proxy) {
     });
 }
 
-// async fn connect_to_app(state: &AppState, app_id: &str, port: u16) -> Result<TcpStream> {
-//     let host = state.lock().select_a_host(app_id).context(format!("tapp {app_id} not found"))?;
-//     TcpStream::connect((host.ip, port))
-// }
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_destination() {
+        let base_domain = ".example.com";
+
+        // Test basic app_id only
+        let result = parse_destination("myapp.example.com", base_domain).unwrap();
+        assert_eq!(result.app_id, "myapp");
+        assert_eq!(result.port, 80);
+        assert!(!result.is_tls);
+
+        // Test app_id with custom port
+        let result = parse_destination("myapp-8080.example.com", base_domain).unwrap();
+        assert_eq!(result.app_id, "myapp");
+        assert_eq!(result.port, 8080);
+        assert!(!result.is_tls);
+
+        // Test app_id with TLS
+        let result = parse_destination("myapp-443s.example.com", base_domain).unwrap();
+        assert_eq!(result.app_id, "myapp");
+        assert_eq!(result.port, 443);
+        assert!(result.is_tls);
+
+        // Test app_id with custom port and TLS
+        let result = parse_destination("myapp-8443s.example.com", base_domain).unwrap();
+        assert_eq!(result.app_id, "myapp");
+        assert_eq!(result.port, 8443);
+        assert!(result.is_tls);
+
+        // Test default port but ends with s
+        let result = parse_destination("myapps.example.com", base_domain).unwrap();
+        assert_eq!(result.app_id, "myapps");
+        assert_eq!(result.port, 80);
+        assert!(!result.is_tls);
+
+        // Test default port but ends with s in port part
+        let result = parse_destination("myapp-s.example.com", base_domain).unwrap();
+        assert_eq!(result.app_id, "myapp");
+        assert_eq!(result.port, 443);
+        assert!(result.is_tls);
+    }
+
+    #[test]
+    fn test_parse_destination_errors() {
+        let base_domain = ".example.com";
+
+        // Test invalid domain suffix
+        assert!(parse_destination("myapp.wrong.com", base_domain).is_err());
+
+        // Test multiple subdomains
+        assert!(parse_destination("invalid.myapp.example.com", base_domain).is_err());
+
+        // Test invalid port format
+        assert!(parse_destination("myapp-65536.example.com", base_domain).is_err());
+        assert!(parse_destination("myapp-abc.example.com", base_domain).is_err());
+
+        // Test too many parts
+        assert!(parse_destination("myapp-8080-extra.example.com", base_domain).is_err());
+
+        // Test empty app_id
+        assert!(parse_destination("-8080.example.com", base_domain).is_err());
+        assert!(parse_destination("myapp-8080ss.example.com", base_domain).is_err());
+    }
+}

tproxy/src/proxy/tls_terminate.rs

@@ -118,9 +118,8 @@ impl TlsTerminateProxy {
         inbound: TcpStream,
         buffer: Vec<u8>,
         app_id: &str,
-        port: Option<u16>,
+        port: u16,
     ) -> Result<()> {
-        let port = port.unwrap_or(80);
         let addresses = self
             .app_state
             .lock()