use async_trait::async_trait;
use bytes::Bytes;
use http_body_util::{BodyExt, Full};
use hyper::{Request, Response, StatusCode};
use std::sync::Arc;
pub type RouteBody = http_body_util::combinators::UnsyncBoxBody<Bytes, hyper::Error>;
#[async_trait]
pub trait RouteHandler: Send + Sync {
async fn handle(&self, req: Request<RouteBody>) -> Response<RouteBody>;
}
pub struct RouteRegistry {
routes: Vec<(String, Arc<dyn RouteHandler>)>,
}
impl RouteRegistry {
pub fn new() -> Self {
Self { routes: Vec::new() }
}
pub fn add_route(&mut self, path: &str, handler: Arc<dyn RouteHandler>) {
self.routes.push((path.to_string(), handler));
}
pub fn is_empty(&self) -> bool {
self.routes.is_empty()
}
pub fn match_route(
&self,
path: &str,
) -> Result<Option<&Arc<dyn RouteHandler>>, RouteValidationError> {
validate_path(path)?;
for (registered_path, handler) in &self.routes {
if path == registered_path {
return Ok(Some(handler));
}
}
Ok(None)
}
}
impl Default for RouteRegistry {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum RouteValidationError {
PathTraversal,
DoubleSlash,
EncodedSeparator,
NullByte,
}
impl std::fmt::Display for RouteValidationError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::PathTraversal => write!(f, "Path traversal detected"),
Self::DoubleSlash => write!(f, "Double slash detected"),
Self::EncodedSeparator => write!(f, "Percent-encoded separator detected"),
Self::NullByte => write!(f, "Null byte detected"),
}
}
}
impl std::error::Error for RouteValidationError {}
impl RouteValidationError {
pub fn into_response(self) -> Response<RouteBody> {
Response::builder()
.status(StatusCode::BAD_REQUEST)
.header("Content-Type", "text/plain")
.body(
Full::new(Bytes::from(format!("Bad Request: {}", self)))
.map_err(|never| match never {})
.boxed_unsync(),
)
.unwrap()
}
}
fn validate_path(path: &str) -> Result<(), RouteValidationError> {
if path.contains('\0') {
return Err(RouteValidationError::NullByte);
}
if path.contains("//") {
return Err(RouteValidationError::DoubleSlash);
}
if path.contains("/../")
|| path.contains("/./")
|| path.ends_with("/..")
|| path.ends_with("/.")
|| path == ".."
|| path == "."
{
return Err(RouteValidationError::PathTraversal);
}
let lower = path.to_ascii_lowercase();
if lower.contains("%2f") || lower.contains("%2e") || lower.contains("%00") {
return Err(RouteValidationError::EncodedSeparator);
}
if lower.contains("%252f") || lower.contains("%252e") {
return Err(RouteValidationError::EncodedSeparator);
}
let bytes = path.as_bytes();
for i in 0..bytes.len() {
if bytes[i] == b'%'
&& (i + 2 >= bytes.len()
|| !bytes[i + 1].is_ascii_hexdigit()
|| !bytes[i + 2].is_ascii_hexdigit())
{
return Err(RouteValidationError::EncodedSeparator);
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use async_trait::async_trait;
struct TestHandler {
body: String,
}
#[async_trait]
impl RouteHandler for TestHandler {
async fn handle(&self, _req: Request<RouteBody>) -> Response<RouteBody> {
Response::builder()
.status(StatusCode::OK)
.body(
Full::new(Bytes::from(self.body.clone()))
.map_err(|never| match never {})
.boxed_unsync(),
)
.unwrap()
}
}
fn registry_with_well_known() -> RouteRegistry {
let mut registry = RouteRegistry::new();
registry.add_route(
"/.well-known/oauth-protected-resource",
Arc::new(TestHandler {
body: r#"{"resource":"https://example.com/mcp"}"#.to_string(),
}),
);
registry
}
#[test]
fn test_route_registry_exact_match() {
let registry = registry_with_well_known();
let result = registry
.match_route("/.well-known/oauth-protected-resource")
.unwrap();
assert!(result.is_some());
}
#[test]
fn test_route_registry_no_prefix_match() {
let registry = registry_with_well_known();
let result = registry
.match_route("/.well-known/oauth-protected-resource/extra")
.unwrap();
assert!(result.is_none());
let result = registry.match_route("/.well-known").unwrap();
assert!(result.is_none());
}
#[test]
fn test_route_registry_case_sensitive() {
let registry = registry_with_well_known();
let result = registry
.match_route("/.Well-Known/OAuth-Protected-Resource")
.unwrap();
assert!(result.is_none());
}
#[test]
fn test_route_registry_reject_path_traversal() {
let registry = registry_with_well_known();
assert!(matches!(
registry.match_route("/../.well-known/oauth-protected-resource"),
Err(RouteValidationError::PathTraversal)
));
assert!(matches!(
registry.match_route("/.well-known/../admin"),
Err(RouteValidationError::PathTraversal)
));
assert!(matches!(
registry.match_route("/./test"),
Err(RouteValidationError::PathTraversal)
));
assert!(matches!(
registry.match_route("/test/.."),
Err(RouteValidationError::PathTraversal)
));
}
#[test]
fn test_route_reject_percent_encoded_slash() {
let registry = registry_with_well_known();
assert!(matches!(
registry.match_route("/.well-known%2foauth-protected-resource"),
Err(RouteValidationError::EncodedSeparator)
));
assert!(matches!(
registry.match_route("/.well-known%2Foauth-protected-resource"),
Err(RouteValidationError::EncodedSeparator)
));
}
#[test]
fn test_route_reject_double_encoding() {
let registry = registry_with_well_known();
assert!(matches!(
registry.match_route("/.well-known%252foauth"),
Err(RouteValidationError::EncodedSeparator)
));
assert!(matches!(
registry.match_route("/.well-known%252etest"),
Err(RouteValidationError::EncodedSeparator)
));
}
#[test]
fn test_route_reject_null_byte() {
let registry = registry_with_well_known();
assert!(matches!(
registry.match_route("/.well-known\x00/test"),
Err(RouteValidationError::NullByte)
));
}
#[test]
fn test_route_reject_double_slash() {
let registry = registry_with_well_known();
assert!(matches!(
registry.match_route("//.well-known/oauth-protected-resource"),
Err(RouteValidationError::DoubleSlash)
));
}
#[test]
fn test_route_no_match_returns_none() {
let registry = registry_with_well_known();
let result = registry.match_route("/not-registered").unwrap();
assert!(result.is_none());
}
#[test]
fn test_route_reject_malformed_percent_encoding() {
let registry = registry_with_well_known();
assert!(matches!(
registry.match_route("/test%"),
Err(RouteValidationError::EncodedSeparator)
));
assert!(matches!(
registry.match_route("/test%2"),
Err(RouteValidationError::EncodedSeparator)
));
assert!(matches!(
registry.match_route("/test%ZZ"),
Err(RouteValidationError::EncodedSeparator)
));
assert!(matches!(
registry.match_route("/test%G1"),
Err(RouteValidationError::EncodedSeparator)
));
}
#[test]
fn test_empty_registry() {
let registry = RouteRegistry::new();
assert!(registry.is_empty());
let result = registry.match_route("/anything").unwrap();
assert!(result.is_none());
}
}