use std::collections::{BTreeMap, BTreeSet};
use std::path::Path;
use anyhow::{bail, Context, Result};
use base64::Engine;
use rand::Rng;
use crate::utils::env::{EnvSource, SystemEnv};
pub const TOKEN_ENV: &str = "OMNI_BRIDGE_TOKEN";
pub const BRIDGE_HEADER: &str = "x-omni-bridge";
pub const BRIDGE_HEADER_VALUE: &str = "1";
pub const BRIDGE_TARGET_HEADER: &str = "x-omni-bridge-target";
const TOKEN_BYTES: usize = 32;
#[must_use]
pub fn generate_token() -> String {
let mut bytes = [0u8; TOKEN_BYTES];
rand::rng().fill_bytes(&mut bytes);
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes)
}
pub fn resolve_token(token_file: Option<&Path>) -> Result<String> {
resolve_token_with(&SystemEnv, token_file)
}
pub(crate) fn resolve_token_with(
env: &impl EnvSource,
token_file: Option<&Path>,
) -> Result<String> {
if let Some(path) = token_file {
return read_token_file(path);
}
if let Some(value) = env.var(TOKEN_ENV) {
let trimmed = value.trim();
if !trimmed.is_empty() {
return Ok(trimmed.to_string());
}
}
Ok(generate_token())
}
pub fn resolve_existing_token(token_file: Option<&Path>) -> Result<String> {
resolve_existing_token_with(&SystemEnv, token_file)
}
pub(crate) fn resolve_existing_token_with(
env: &impl EnvSource,
token_file: Option<&Path>,
) -> Result<String> {
if let Some(path) = token_file {
return read_token_file(path);
}
match env.var(TOKEN_ENV) {
Some(value) if !value.trim().is_empty() => Ok(value.trim().to_string()),
_ => bail!(
"No session token found. Set {TOKEN_ENV} or pass --token-file with the token the \
running bridge printed."
),
}
}
fn read_token_file(path: &Path) -> Result<String> {
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let meta = std::fs::metadata(path)
.with_context(|| format!("Failed to stat token file {}", path.display()))?;
let mode = meta.permissions().mode() & 0o777;
if mode & 0o077 != 0 {
bail!(
"Token file {} must be 0600 (owner-only); found {:o}",
path.display(),
mode
);
}
}
let contents = std::fs::read_to_string(path)
.with_context(|| format!("Failed to read token file {}", path.display()))?;
let trimmed = contents.trim();
if trimmed.is_empty() {
bail!("Token file {} is empty", path.display());
}
Ok(trimmed.to_string())
}
#[must_use]
pub fn constant_time_eq(a: &str, b: &str) -> bool {
let (a, b) = (a.as_bytes(), b.as_bytes());
if a.len() != b.len() {
return false;
}
let mut diff = 0u8;
for (x, y) in a.iter().zip(b.iter()) {
diff |= x ^ y;
}
diff == 0
}
#[must_use]
pub fn bearer_matches(authorization: Option<&str>, token: &str) -> bool {
let Some(value) = authorization else {
return false;
};
let Some(presented) = value.strip_prefix("Bearer ") else {
return false;
};
constant_time_eq(presented.trim(), token)
}
#[must_use]
pub fn has_bridge_header(value: Option<&str>) -> bool {
value.is_some_and(|v| v.trim() == BRIDGE_HEADER_VALUE)
}
#[must_use]
pub fn host_allowed(host: &str, control_port: u16) -> bool {
let allowed = [
format!("localhost:{control_port}"),
format!("127.0.0.1:{control_port}"),
format!("[::1]:{control_port}"),
];
allowed.iter().any(|a| a == host)
}
#[must_use]
pub fn is_browser_originated(origin: Option<&str>, sec_fetch_site: Option<&str>) -> bool {
if origin.is_some() {
return true;
}
matches!(
sec_fetch_site.map(str::trim),
Some("cross-site" | "same-site" | "same-origin")
)
}
#[must_use]
pub fn header_is_safe(name: &str, value: &str) -> bool {
let bad = |s: &str| s.bytes().any(|b| b == b'\r' || b == b'\n' || b == 0);
!name.is_empty() && !bad(name) && !bad(value)
}
#[must_use]
pub fn normalize_request_path(raw: &str) -> Option<String> {
let decoded = percent_decode(raw)?;
if decoded.bytes().any(|b| b == b'\r' || b == b'\n' || b == 0) {
return None;
}
if decoded
.split('/')
.any(|seg| seg == ".." || seg == "..%2f" || seg == "..%2F")
{
return None;
}
Some(decoded)
}
fn percent_decode(raw: &str) -> Option<String> {
let bytes = raw.as_bytes();
let mut out = Vec::with_capacity(bytes.len());
let mut i = 0;
while i < bytes.len() {
match bytes[i] {
b'%' => {
let hi = bytes.get(i + 1).copied().and_then(hex_val)?;
let lo = bytes.get(i + 2).copied().and_then(hex_val)?;
out.push(hi << 4 | lo);
i += 3;
}
b => {
out.push(b);
i += 1;
}
}
}
String::from_utf8(out).ok()
}
fn hex_val(b: u8) -> Option<u8> {
match b {
b'0'..=b'9' => Some(b - b'0'),
b'a'..=b'f' => Some(b - b'a' + 10),
b'A'..=b'F' => Some(b - b'A' + 10),
_ => None,
}
}
#[derive(Debug, PartialEq, Eq)]
pub enum ScopeError {
CrossOriginDenied,
Malformed,
}
pub fn validate_outbound_url(url: &str, allowed: &[&str]) -> Result<(), ScopeError> {
let is_relative = url.starts_with('/') && !url.starts_with("//");
if is_relative {
if url.bytes().any(|b| b == b'\r' || b == b'\n' || b == 0) {
return Err(ScopeError::Malformed);
}
return Ok(());
}
if allowed.is_empty() {
return Err(ScopeError::CrossOriginDenied);
}
let target = url::Url::parse(url).map_err(|_| ScopeError::Malformed)?;
let permitted = allowed
.iter()
.filter_map(|a| url::Url::parse(a).ok())
.any(|allow| origins_match(&target, &allow));
if permitted {
Ok(())
} else {
Err(ScopeError::CrossOriginDenied)
}
}
fn origins_match(a: &url::Url, b: &url::Url) -> bool {
a.scheme() == b.scheme()
&& a.host_str() == b.host_str()
&& a.port_or_known_default() == b.port_or_known_default()
}
fn canonical_origin(s: &str) -> Option<String> {
let origin = url::Url::parse(s.trim()).ok()?.origin();
origin.is_tuple().then(|| origin.ascii_serialization())
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct OriginAllowlist {
map: BTreeMap<String, BTreeSet<String>>,
}
impl OriginAllowlist {
pub fn parse<S: AsRef<str>>(values: &[S]) -> Result<Self, String> {
let mut map: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();
for raw in values {
let raw = raw.as_ref().trim();
if raw.is_empty() {
return Err("empty --allow-origin value".to_string());
}
let (connect, outbound) = match raw.split_once('=') {
Some((c, o)) => (c.trim(), o.trim()),
None => (raw, raw),
};
let connect = canonical_origin(connect).ok_or_else(|| {
format!("invalid --allow-origin connecting origin: {connect:?} (in {raw:?})")
})?;
let outbound = canonical_origin(outbound).ok_or_else(|| {
format!("invalid --allow-origin outbound origin: {outbound:?} (in {raw:?})")
})?;
map.entry(connect).or_default().insert(outbound);
}
Ok(Self { map })
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
#[must_use]
pub fn permits_connection(&self, origin: Option<&str>) -> bool {
if self.map.is_empty() {
return true;
}
origin
.and_then(canonical_origin)
.is_some_and(|o| self.map.contains_key(&o))
}
#[must_use]
pub fn outbound_for(&self, origin: Option<&str>) -> Vec<&str> {
origin
.and_then(canonical_origin)
.and_then(|o| self.map.get(&o))
.map(|set| set.iter().map(String::as_str).collect())
.unwrap_or_default()
}
#[must_use]
pub fn describe(&self) -> Vec<String> {
self.map
.iter()
.map(|(connect, outbound)| {
let outbound = outbound.iter().cloned().collect::<Vec<_>>().join(", ");
format!("{connect} → {outbound}")
})
.collect()
}
}
#[must_use]
pub fn ws_subprotocol_token<'a, I>(subprotocols: I, token: &str) -> Option<&'a str>
where
I: IntoIterator<Item = &'a str>,
{
subprotocols
.into_iter()
.map(str::trim)
.find(|p| constant_time_eq(p, token))
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
#[test]
fn generated_tokens_are_unique_and_urlsafe() {
let a = generate_token();
let b = generate_token();
assert_ne!(a, b);
assert!(a
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_'));
assert!(a.len() >= 40);
}
#[test]
fn constant_time_eq_matches_str_eq() {
assert!(constant_time_eq("abc", "abc"));
assert!(!constant_time_eq("abc", "abd"));
assert!(!constant_time_eq("abc", "abcd"));
}
#[test]
fn bearer_accepts_only_correct_token() {
assert!(bearer_matches(Some("Bearer tok"), "tok"));
assert!(!bearer_matches(Some("Bearer wrong"), "tok"));
assert!(!bearer_matches(Some("tok"), "tok"));
assert!(!bearer_matches(None, "tok"));
}
#[test]
fn bridge_header_must_be_one() {
assert!(has_bridge_header(Some("1")));
assert!(has_bridge_header(Some(" 1 ")));
assert!(!has_bridge_header(Some("0")));
assert!(!has_bridge_header(None));
}
#[test]
fn host_allowlist_blocks_rebinding() {
assert!(host_allowed("localhost:9998", 9998));
assert!(host_allowed("127.0.0.1:9998", 9998));
assert!(host_allowed("[::1]:9998", 9998));
assert!(!host_allowed("evil.example.com:9998", 9998));
assert!(!host_allowed("localhost:9999", 9998));
assert!(!host_allowed("localhost", 9998));
}
#[test]
fn browser_origin_is_rejected() {
assert!(is_browser_originated(Some("https://evil.test"), None));
assert!(is_browser_originated(None, Some("cross-site")));
assert!(is_browser_originated(None, Some("same-site")));
assert!(!is_browser_originated(None, None));
assert!(!is_browser_originated(None, Some("none")));
}
#[test]
fn header_crlf_is_rejected() {
assert!(header_is_safe("Accept", "application/json"));
assert!(!header_is_safe("X\r\nEvil", "v"));
assert!(!header_is_safe("X", "a\r\nSet-Cookie: y"));
assert!(!header_is_safe("", "v"));
}
#[test]
fn path_normalization_rejects_traversal() {
assert_eq!(
normalize_request_path("/loki/api/v1/labels").as_deref(),
Some("/loki/api/v1/labels")
);
assert_eq!(
normalize_request_path("/a/%2e%2e/b"),
Some("/a/../b".to_string()).filter(|_| false).or(None)
);
assert!(normalize_request_path("/a/../b").is_none());
assert!(normalize_request_path("/a/%2e%2e/b").is_none());
assert!(normalize_request_path("/a/%00/b").is_none());
assert!(normalize_request_path("/bad%2").is_none());
}
#[test]
fn outbound_scope_is_default_closed() {
assert_eq!(validate_outbound_url("/api/foo", &[]), Ok(()));
assert_eq!(
validate_outbound_url("https://evil.test/x", &[]),
Err(ScopeError::CrossOriginDenied)
);
assert_eq!(
validate_outbound_url("//evil.test/x", &[]),
Err(ScopeError::CrossOriginDenied)
);
}
#[test]
fn outbound_scope_honors_allowed_origins() {
assert_eq!(
validate_outbound_url("https://ok.test/x", &["https://ok.test"]),
Ok(())
);
assert_eq!(
validate_outbound_url("https://evil.test/x", &["https://ok.test"]),
Err(ScopeError::CrossOriginDenied)
);
assert_eq!(
validate_outbound_url("https://b.test/x", &["https://a.test", "https://b.test"]),
Ok(())
);
assert_eq!(validate_outbound_url("/x", &["https://ok.test"]), Ok(()));
}
#[test]
fn ws_subprotocol_token_selects_match() {
assert_eq!(ws_subprotocol_token(["a", "tok", "b"], "tok"), Some("tok"));
assert_eq!(ws_subprotocol_token(["a", "b"], "tok"), None);
assert_eq!(ws_subprotocol_token([], "tok"), None);
}
#[test]
fn empty_allowlist_opens_the_ws_gate() {
let empty = OriginAllowlist::default();
assert!(empty.is_empty());
assert!(empty.permits_connection(Some("https://anything.test")));
assert!(empty.permits_connection(None));
assert!(empty.outbound_for(Some("https://anything.test")).is_empty());
}
#[test]
fn allowlist_gates_connection_by_configured_key() {
let list = OriginAllowlist::parse(&["https://ok.test"]).unwrap();
assert!(list.permits_connection(Some("https://ok.test")));
assert!(list.permits_connection(Some("https://ok.test:443")));
assert!(!list.permits_connection(Some("https://evil.test")));
assert!(!list.permits_connection(None));
}
#[test]
fn shorthand_grants_a_tab_its_own_origin() {
let list = OriginAllowlist::parse(&["https://grafana.internal"]).unwrap();
assert_eq!(
list.outbound_for(Some("https://grafana.internal")),
vec!["https://grafana.internal"]
);
assert!(list.outbound_for(Some("https://other.test")).is_empty());
assert!(list.outbound_for(None).is_empty());
}
#[test]
fn mapping_scopes_outbound_per_connecting_origin() {
let list = OriginAllowlist::parse(&[
"https://grafana.internal",
"https://www.facebook.com=https://static.xx.fbcdn.net",
])
.unwrap();
assert_eq!(
list.outbound_for(Some("https://grafana.internal")),
vec!["https://grafana.internal"]
);
assert_eq!(
list.outbound_for(Some("https://www.facebook.com")),
vec!["https://static.xx.fbcdn.net"]
);
assert_eq!(
validate_outbound_url(
"https://static.xx.fbcdn.net/x",
&list.outbound_for(Some("https://grafana.internal"))
),
Err(ScopeError::CrossOriginDenied)
);
}
#[test]
fn repeated_keys_accumulate_outbound_origins() {
let list = OriginAllowlist::parse(&[
"https://app.test=https://a.cdn.test",
"https://app.test=https://b.cdn.test",
])
.unwrap();
assert_eq!(
list.outbound_for(Some("https://app.test")),
vec!["https://a.cdn.test", "https://b.cdn.test"]
);
}
#[test]
fn parse_rejects_malformed_and_empty_values() {
assert!(OriginAllowlist::parse(&["not a url"]).is_err());
assert!(OriginAllowlist::parse(&["https://ok.test=nonsense"]).is_err());
assert!(OriginAllowlist::parse(&[""]).is_err());
assert!(OriginAllowlist::parse(&["=https://ok.test"]).is_err());
}
#[cfg(unix)]
#[test]
fn token_file_requires_0600() {
use std::io::Write;
use std::os::unix::fs::PermissionsExt;
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("tok");
let mut f = std::fs::File::create(&path).unwrap();
writeln!(f, "secret-token").unwrap();
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o644)).unwrap();
assert!(resolve_token(Some(&path)).is_err());
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600)).unwrap();
assert_eq!(resolve_token(Some(&path)).unwrap(), "secret-token");
}
use crate::test_support::env::MapEnv;
#[test]
fn resolve_token_reads_trimmed_env_var() {
let env = MapEnv::new().with(TOKEN_ENV, " env-token ");
assert_eq!(resolve_token_with(&env, None).unwrap(), "env-token");
}
#[test]
fn resolve_token_generates_when_env_empty_or_absent() {
let a = resolve_token_with(&MapEnv::new(), None).unwrap();
assert!(a.len() >= 40);
let env = MapEnv::new().with(TOKEN_ENV, " ");
let b = resolve_token_with(&env, None).unwrap();
assert!(b.len() >= 40);
assert_ne!(a, b);
}
#[test]
fn resolve_existing_token_reads_env_var() {
let env = MapEnv::new().with(TOKEN_ENV, "client-token");
assert_eq!(
resolve_existing_token_with(&env, None).unwrap(),
"client-token"
);
}
#[test]
fn resolve_existing_token_errors_without_source() {
let err = resolve_existing_token_with(&MapEnv::new(), None).unwrap_err();
assert!(err.to_string().contains(TOKEN_ENV));
}
#[test]
fn resolve_existing_token_reads_file() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("tok");
std::fs::write(&path, " file-token\n").unwrap();
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600)).unwrap();
}
assert_eq!(resolve_existing_token(Some(&path)).unwrap(), "file-token");
}
#[test]
fn token_file_missing_errors() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("does-not-exist");
assert!(resolve_token(Some(&path)).is_err());
}
#[test]
fn token_file_empty_errors() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("empty");
std::fs::write(&path, " \n").unwrap();
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600)).unwrap();
}
let err = resolve_token(Some(&path)).unwrap_err();
assert!(err.to_string().contains("empty"));
}
}