Allow restrictions based on custom Authorization Header matching logic

restrictions.yaml

@@ -10,9 +10,15 @@ restrictions:
       # The regex does a match, so if you want to match exactly you need to bound the pattern with ^ $
       # I.e: "tesotron" is going to match "XXXtesotronXXX", but "^tesotron$" is going to match only "tesotron"
       - !PathPrefix "^.*$"
-      # This match applies only if it succeeds to match the Authentication Header with the given regex.
-      # If present, Authentication Header must exists and must match the regex.
+      # This match applies only if it succeeds to match the Authorization Header with the given regex.
+      # If present, Authorization Header must exists and must match the regex.
       # - !Authorization "^[Bb]earer +actual_bearer_token_to_match$"
+      # This match applies if the Authorization Header is allowed according to a custom lua script.
+      # The lua script must contain a global function named 'auth_validate' that gets the Auth Header value
+      # as a string parameter, and returns true iff the custom logic decides the value is authorized.
+      # Any failure to load/execute the lua script is logged, and considered authorization failure.
+      # If the match is present but no Auth Header exists, it's considered authorization failure.
+      # - !AuthorizationScript "/etc/wstunnel/authrorize.lua"
       # The only other possible match type for now is !Any, that match everything/any request
       # - !Any
 

wstunnel/Cargo.toml

@@ -56,6 +56,8 @@ rcgen = { version = "0.13.2", default-features = false, features = [] }
 hickory-resolver = { version = "0.25.2", default-features = false, features = ["system-config", "tokio", "rustls-platform-verifier"] }
 aws-lc-rs = { version = "*", optional = true }
 
+mlua = { version = "0.10.5", features = ["lua54", "vendored"] }
+
 [target.'cfg(not(target_family = "unix"))'.dependencies]
 crossterm = { version = "0.29.0" }
 tokio-util = { version = "0.7.15", features = ["io"] }

wstunnel/src/restrictions/types.rs

@@ -25,6 +25,7 @@ pub enum MatchConfig {
     PathPrefix(Regex),
     #[serde(with = "serde_regex")]
     Authorization(Regex),
+    AuthorizationScript(String),
 }
 
 #[derive(Debug, Clone, Deserialize)]

wstunnel/src/tunnel/server/utils.rs

@@ -12,6 +12,7 @@ use hyper::body::{Body, Incoming};
 use hyper::header::{AUTHORIZATION, COOKIE, HeaderValue, SEC_WEBSOCKET_PROTOCOL};
 use hyper::{Request, Response, StatusCode, http};
 use jsonwebtoken::TokenData;
+use mlua::prelude::{Lua, LuaFunction};
 use std::net::IpAddr;
 use tracing::{error, info, warn};
 use url::Host;
@@ -116,10 +117,47 @@ impl RestrictionConfig {
             MatchConfig::Any => true,
             MatchConfig::PathPrefix(path) => path.is_match(path_prefix),
             MatchConfig::Authorization(auth) => authorization_header_val.is_some_and(|val| auth.is_match(val)),
+            MatchConfig::AuthorizationScript(script_name) => {
+                authorization_header_val.is_some_and(|val| auth_header_matcher(script_name, val))
+            }
         })
     }
 }
 
+fn auth_header_matcher(script_name: &String, auth_val: &str) -> bool {
+    let validate_fn_name = "auth_validate";
+    let lua = Lua::new();
+
+    let script = match std::fs::read_to_string(script_name) {
+        Ok(s) => s,
+        Err(e) => {
+            error!("Failed to read {}: {}", script_name, e);
+            return false;
+        }
+    };
+
+    if let Err(e) = lua.load(&script).exec() {
+        error!("Failed to load lua script {}: {}", script_name, e);
+        return false;
+    }
+
+    let validate_fn: LuaFunction = match lua.globals().get(validate_fn_name) {
+        Ok(func) => func,
+        Err(e) => {
+            error!("Failed to find '{}' lua function: {}", validate_fn_name, e);
+            return false;
+        }
+    };
+
+    match validate_fn.call(auth_val) {
+        Ok(result) => result,
+        Err(e) => {
+            error!("Failed calling '{}' lua function: {}", validate_fn_name, e);
+            false
+        }
+    }
+}
+
 impl AllowReverseTunnelConfig {
     #[inline]
     fn is_allowed(&self, remote: &RemoteAddr) -> bool {