feat: location host match supports regex

src/proxy/location.rs

@@ -85,12 +85,56 @@ fn new_path_selector(path: &str) -> Result<PathSelector> {
     Ok(se)
 }
 
+#[derive(Debug)]
+struct RegexHost {
+    value: Regex,
+}
+
+#[derive(Debug)]
+struct EqualHost {
+    value: String,
+}
+
+#[derive(Debug)]
+enum HostSelector {
+    RegexHost(RegexHost),
+    EqualHost(EqualHost),
+}
+
+/// New a host selector, regex or  equal selector
+fn new_host_selector(host: &str) -> Result<HostSelector> {
+    let host = host.trim();
+    if host.is_empty() {
+        return Ok(HostSelector::EqualHost(EqualHost {
+            value: host.to_string(),
+        }));
+    }
+    let first = host.chars().next().unwrap_or_default();
+    let last = host.substring(1, host.len()).trim();
+    let se = match first {
+        '~' => {
+            let re = Regex::new(last).context(RegexSnafu {
+                value: last.to_string(),
+            })?;
+            HostSelector::RegexHost(RegexHost { value: re })
+        },
+        _ => {
+            // trim
+            HostSelector::EqualHost(EqualHost {
+                value: host.trim().to_string(),
+            })
+        },
+    };
+
+    Ok(se)
+}
+
 pub struct Location {
     pub name: String,
     pub key: String,
     path: String,
     path_selector: PathSelector,
-    hosts: Vec<String>,
+    hosts: Vec<HostSelector>,
     reg_rewrite: Option<(Regex, String)>,
     proxy_add_headers: Option<Vec<HttpHeader>>,
     proxy_set_headers: Option<Vec<HttpHeader>>,
@@ -149,10 +193,19 @@ impl Location {
         let mut hosts = vec![];
         for item in conf.host.clone().unwrap_or_default().split(',') {
             let host = item.trim().to_string();
-            if !host.is_empty() {
-                hosts.push(host);
+            if host.is_empty() {
+                continue;
             }
+            hosts.push(new_host_selector(&host)?);
         }
+        // sort hosts regexp --> equal
+        hosts.sort_by_key(|host| match host {
+            HostSelector::RegexHost(RegexHost { value }) => {
+                1000 + value.to_string().len()
+            },
+            HostSelector::EqualHost(EqualHost { value }) => value.len(),
+        });
+        hosts.reverse();
 
         let path = conf.path.clone().unwrap_or_default();
 
@@ -201,7 +254,14 @@ impl Location {
             return true;
         }
 
-        self.hosts.iter().any(|item| item == host)
+        self.hosts.iter().any(|item| match item {
+            HostSelector::RegexHost(RegexHost { value }) => {
+                let (matched, _) =
+                    util::regex_capture(value, host).unwrap_or_default();
+                matched
+            },
+            HostSelector::EqualHost(EqualHost { value }) => value == host,
+        })
     }
     /// Sets the maximum allowed size of the client request body.
     /// If the size in a request exceeds the configured value, the 413 (Request Entity Too Large) error

src/util/mod.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use ahash::AHashMap;
 use base64::{engine::general_purpose::STANDARD, Engine};
 use bytes::BytesMut;
 use http::{HeaderName, Uri};
@@ -20,6 +21,7 @@ use path_absolutize::*;
 use pingora::cache::CacheKey;
 use pingora::tls::ssl::SslVersion;
 use pingora::{http::RequestHeader, proxy::Session};
+use regex::Regex;
 use snafu::Snafu;
 use std::collections::HashMap;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
@@ -335,6 +337,37 @@ pub fn base64_decode<T: AsRef<[u8]>>(
     STANDARD.decode(data)
 }
 
+pub fn regex_capture(
+    re: &Regex,
+    value: &str,
+) -> Result<(bool, Option<AHashMap<String, String>>), Error> {
+    if !re.is_match(value) {
+        return Ok((false, None));
+    }
+    let mut m = AHashMap::new();
+    let mut keys = vec![];
+    for name in re.capture_names() {
+        keys.push(name.unwrap_or_default());
+    }
+    let Some(cap) = re.captures(value) else {
+        return Ok((true, Some(m)));
+    };
+    for (index, value) in cap.iter().enumerate() {
+        if index >= keys.len() {
+            continue;
+        }
+        let key = keys[index];
+        if key.is_empty() {
+            continue;
+        }
+        let Some(value) = value else {
+            continue;
+        };
+        m.insert(key.to_string(), value.as_str().to_string());
+    }
+    Ok((true, Some(m)))
+}
+
 const B_100: usize = 100;
 const KB: usize = 1_000;
 const KB_100: usize = 100 * KB;