gw: Add flag p in sni

Author: kvinwang

gateway/src/proxy.rs

@@ -70,6 +70,79 @@ struct DstInfo {
     port: u16,
     is_tls: bool,
     is_h2: bool,
+    is_pp: bool,
+}
+
+fn parse_app_addr(addr: &str) -> Result<DstInfo> {
+    let (app_id, port_part) = addr
+        .rsplit_once('-')
+        .or_else(|| addr.rsplit_once(':'))
+        .unwrap_or((addr, ""));
+    if app_id.is_empty() {
+        bail!("app id is empty");
+    }
+    let mut dst = DstInfo {
+        app_id: app_id.to_owned(),
+        port: 80,
+        is_tls: false,
+        is_h2: false,
+        is_pp: false,
+    };
+
+    if port_part.is_empty() {
+        return Ok(dst);
+    };
+
+    // Parse suffixes from right to left: g, s, p
+    let part_bytes = port_part.as_bytes();
+    let mut end_idx = part_bytes.len();
+
+    // Parse from right to left until we hit a digit
+    while end_idx > 0 {
+        let ch = part_bytes[end_idx - 1] as char;
+        match ch {
+            c if c.is_ascii_digit() => {
+                break;
+            }
+            'g' => {
+                if dst.is_h2 {
+                    bail!("invalid app address: duplicate suffix 'g'");
+                }
+                dst.is_h2 = true;
+                end_idx -= 1;
+            }
+            's' => {
+                if dst.is_tls {
+                    bail!("invalid app address: duplicate suffix 's'");
+                }
+                dst.is_tls = true;
+                end_idx -= 1;
+            }
+            'p' => {
+                if dst.is_pp {
+                    bail!("invalid app address: duplicate suffix 'p'");
+                }
+                dst.is_pp = true;
+                end_idx -= 1;
+            }
+            _ => {
+                bail!("invalid app address: unrecognized suffix character '{ch}'");
+            }
+        }
+    }
+
+    if dst.is_h2 && dst.is_tls {
+        bail!("invalid app address: both 's' and 'g' suffixes are present");
+    }
+
+    let port_str = &port_part[..end_idx];
+    let port = if port_str.is_empty() {
+        None
+    } else {
+        Some(port_str.parse::<u16>().context("invalid port")?)
+    };
+    dst.port = port.unwrap_or(if dst.is_tls { 443 } else { 80 });
+    Ok(dst)
 }
 
 fn parse_destination(sni: &str, dotted_base_domain: &str) -> Result<DstInfo> {
@@ -80,53 +153,7 @@ fn parse_destination(sni: &str, dotted_base_domain: &str) -> Result<DstInfo> {
     if subdomain.contains('.') {
         bail!("only one level of subdomain is supported, got sni={sni}, subdomain={subdomain}");
     }
-    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 port;
-    let is_h2;
-    match last_part {
-        None => {
-            is_tls = false;
-            is_h2 = false;
-            port = None;
-        }
-        Some(last_part) => {
-            let (port_str, has_g) = match last_part.strip_suffix('g') {
-                Some(without_g) => (without_g, true),
-                None => (last_part, false),
-            };
-
-            let (port_str, has_s) = match port_str.strip_suffix('s') {
-                Some(without_s) => (without_s, true),
-                None => (port_str, false),
-            };
-            if has_g && has_s {
-                bail!("invalid sni format: `gs` is not allowed");
-            }
-            is_h2 = has_g;
-            is_tls = has_s;
-            port = if port_str.is_empty() {
-                None
-            } else {
-                Some(port_str.parse::<u16>().context("invalid port")?)
-            };
-        }
-    };
-    let port = port.unwrap_or(if is_tls { 443 } else { 80 });
-    if parts.next().is_some() {
-        bail!("invalid sni format");
-    }
-    Ok(DstInfo {
-        app_id,
-        port,
-        is_tls,
-        is_h2,
-    })
+    parse_app_addr(subdomain)
 }
 
 pub static NUM_CONNECTIONS: AtomicU64 = AtomicU64::new(0);

gateway/src/proxy/tls_passthough.rs

@@ -3,37 +3,23 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use anyhow::{Context, Result};
-use std::fmt::Debug;
 use tokio::{io::AsyncWriteExt, net::TcpStream, task::JoinSet, time::timeout};
 use tracing::{debug, info};
 
 use crate::{
     main_service::Proxy,
     models::{Counting, EnteredCounter},
+    proxy::{parse_app_addr, DstInfo},
 };
 
 use super::{io_bridge::bridge, AddressGroup};
 
-#[derive(Debug)]
-struct AppAddress {
-    app_id: String,
-    port: u16,
-}
-
-impl AppAddress {
-    fn parse(data: &[u8]) -> Result<Self> {
-        // format: "3327603e03f5bd1f830812ca4a789277fc31f577:555"
-        let data = String::from_utf8(data.to_vec()).context("invalid app address")?;
-        let (app_id, port) = data.split_once(':').context("invalid app address")?;
-        Ok(Self {
-            app_id: app_id.to_string(),
-            port: port.parse().context("invalid port")?,
-        })
-    }
+fn parse_txt_addr(addr: &[u8]) -> Result<DstInfo> {
+    parse_app_addr(core::str::from_utf8(addr).context("invalid app address")?)
 }
 
 /// resolve app address by sni
-async fn resolve_app_address(prefix: &str, sni: &str, compat: bool) -> Result<AppAddress> {
+async fn resolve_app_address(prefix: &str, sni: &str, compat: bool) -> Result<DstInfo> {
     let txt_domain = format!("{prefix}.{sni}");
     let resolver = hickory_resolver::AsyncResolver::tokio_from_system_conf()
         .context("failed to create dns resolver")?;
@@ -54,7 +40,7 @@ async fn resolve_app_address(prefix: &str, sni: &str, compat: bool) -> Result<Ap
             let Some(data) = txt_record.txt_data().first() else {
                 continue;
             };
-            return AppAddress::parse(data).context("failed to parse app address");
+            return parse_txt_addr(data).context("failed to parse app address");
         }
         anyhow::bail!("failed to resolve app address");
     } else {
@@ -67,7 +53,7 @@ async fn resolve_app_address(prefix: &str, sni: &str, compat: bool) -> Result<Ap
             .txt_data()
             .first()
             .context("no data in txt record")?;
-        AppAddress::parse(data).context("failed to parse app address")
+        parse_txt_addr(data).context("failed to parse app address")
     }
 }