alkdev/alkhttp
6
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(server): single token resolution on WS upgrade and /mcp routes (SRV-11)

Browse Source 
The router-wide bearer_auth_middleware route_layer wraps every route
registered before the call (axum 0.8 RouteManager semantics, verified
against the vendored axum 0.8.9 source), so it double-wrapped the WS
upgrade route (own ws_bearer_auth layer) and the /mcp nest (own inner
bearer layer) — the token resolved twice per request on both, and the
SRV-10 comment claimed the opposite.

Restructure build_router: /mcp is merged and the WS upgrade route is
registered after the router-wide route_layer, so each keeps exactly one
auth layer. The WS MethodRouter now carries the decoy 405 fallback
explicitly (MethodRouter::route_layer wraps method endpoints, not the
fallback), preserving the SRV-07 decoy shape for wrong-method probes on
/alk/channels. Comments state axum's actual route_layer semantics.

Tests: a counting IdentityProvider pins one resolution per WS upgrade
request and per /mcp initialize (verified to fail with left: 2/3 under
the pre-fix ordering); the 401 enforcement and the WS-path decoy 405
are pinned.

Verification: cargo test (304), cargo test --all-features (376 + 5
integration suites), clippy -D warnings (default + all-features),
cargo fmt --check, cargo doc --no-deps
This commit is contained in:
glm-5.3-flash
2026-08-30 12:17:32 +00:00
parent 6cd0fd3f86
commit f431dea5b5
1 changed files with 222 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/server/adapter.rs
+222 -29
View File
@@ -300,27 +300,8 @@ fn build_router(state: RouterState, extra_routes: Option<Router>) -> Router {
// cache (SRV-09), threaded through `RouterState` below.
.route("/openapi.json", get(openapi_json_handler))
.route("/healthz", get(healthz))
// The WS upgrade route carries its own bearer middleware
// (`ws_bearer_auth` — 401 without a resolvable token, because a WS
// session without an identity cannot run AccessControl::check).
// This route is registered BEFORE the router-wide
// bearer_auth_middleware route_layer below (route_layer applies
// only to routes registered before it), so the WS path resolves
// the token exactly once — enforced (401) — while the gateway
// endpoints resolve once, permissively (no enforcement; the
// dispatch's AccessControl decides). See SRV-10: the previous
// mux nested this route inside the shared layer and resolved
// the same token twice.
.route(
WS_UPGRADE_PATH,
get(crate::websocket::ws_upgrade_handler).route_layer(from_fn_with_state(
auth_state.clone(),
crate::websocket::ws_bearer_auth,
)),
)
.fallback(decoy_fallback)
.method_not_allowed_fallback(decoy_method_not_allowed)
.merge(mcp_router);
.method_not_allowed_fallback(decoy_method_not_allowed);
let with_extras = match extra_routes {
Some(extra) => {
@@ -339,16 +320,48 @@ fn build_router(state: RouterState, extra_routes: Option<Router>) -> Router {
let with_extras = with_extras.method_not_allowed_fallback(decoy_method_not_allowed);
// Applied after the merges (ADR-046 §4): the bearer-auth layer wraps
// the extra routes and every default-surface route registered above
// except the WS upgrade route (registered earlier with its own
// enforced layer, so the token resolves once per request — SRV-10).
// A route that wants to opt out carries its own inner auth layer,
// which runs innermost (last-applied wins for extension stashes it
// inserts).
let with_extras =
with_extras.route_layer(from_fn_with_state(auth_state, bearer_auth_middleware));
// every route registered before this call (the gateway endpoints,
// /openapi.json, /healthz, and the extra routes) — axum 0.8
// semantics are that `route_layer` wraps the routes registered
// before it, not the ones after. Routes that must not pass through
// this layer (the /mcp nest and the WS upgrade route, each with its
// own auth) are merged/registered after it — SRV-11: before the
// reorder, a route carrying an inner auth layer was also wrapped by
// this one and resolved the token twice (the SRV-10 double-resolve,
// the old comment claiming the opposite).
let with_auth = with_extras.route_layer(from_fn_with_state(
Arc::clone(&auth_state),
bearer_auth_middleware,
));
with_extras.with_state(state)
// Merged after the router-wide route_layer on purpose: the /mcp
// nest carries its own bearer layer (applied around the nested
// service, `from_fn_with_state` above), so registering it here keeps
// exactly one token resolution per request (SRV-11). nest_service
// registers a plain Route endpoint (no MethodRouter), so the decoy
// 405 fallback has no interplay with this merge.
let with_mcp = with_auth.merge(mcp_router);
// Registered after the router-wide route_layer on purpose: axum's
// route_layer applies only to earlier-registered routes, so the WS
// upgrade path resolves the token exactly once through its own
// `ws_bearer_auth` (401 without a resolvable token — a WS session
// without an identity cannot run AccessControl::check). The
// MethodRouter carries the decoy 405 fallback explicitly
// (MethodRouter::route_layer wraps method endpoints, not the
// fallback) — a wrong-method probe on /alk/channels keeps the decoy
// shape instead of axum's bare 405.
with_mcp
.route(
WS_UPGRADE_PATH,
get(crate::websocket::ws_upgrade_handler)
.fallback(decoy_method_not_allowed)
.route_layer(from_fn_with_state(
Arc::clone(&auth_state),
crate::websocket::ws_bearer_auth,
)),
)
.with_state(state)
}
/// Enforce the per-method reserved-path collision rule (ADR-046 §3):
@@ -1043,6 +1056,65 @@ mod tests {
Arc::new(StaticProvider)
}
struct CountingProvider {
resolutions: std::sync::Mutex<usize>,
}
impl CountingProvider {
fn new() -> Self {
Self {
resolutions: std::sync::Mutex::new(0),
}
}
fn resolutions(&self) -> usize {
*self.resolutions.lock().unwrap_or_else(|e| e.into_inner())
}
}
impl IdentityProvider for CountingProvider {
fn resolve_from_fingerprint(&self, _: &str) -> Option<alkcall::core::auth::Identity> {
None
}
fn resolve_from_token(
&self,
_: &alkcall::core::auth::AuthToken,
) -> Option<alkcall::core::auth::Identity> {
*self.resolutions.lock().unwrap_or_else(|e| e.into_inner()) += 1;
Some(alkcall::core::auth::Identity {
id: "worker-a".to_string(),
scopes: vec![],
resources: std::collections::HashMap::new(),
})
}
}
async fn ws_upgrade_oneshot(
app: Router,
authorization: &str,
) -> axum::http::Response<axum::body::Body> {
use tower::ServiceExt;
let request = axum::http::Request::builder()
.method(axum::http::Method::GET)
.uri(WS_UPGRADE_PATH)
.header(axum::http::header::AUTHORIZATION, authorization)
.header(axum::http::header::CONNECTION, "upgrade")
.header(axum::http::header::UPGRADE, "websocket")
.header(axum::http::header::SEC_WEBSOCKET_VERSION, "13")
.header(
axum::http::header::SEC_WEBSOCKET_KEY,
"dGhlIHNhbXBsZSBub25jZQ==",
)
.header(
axum::http::header::HOST,
axum::http::HeaderValue::from_static("localhost"),
);
let mut request = request.body(axum::body::Body::empty()).unwrap();
let on_upgrade = hyper::upgrade::on(&mut request);
request.extensions_mut().insert(on_upgrade);
app.oneshot(request).await.unwrap()
}
fn router_state(idp: Arc<dyn IdentityProvider>) -> RouterState {
RouterState {
registry: empty_registry(),
@@ -1290,6 +1362,127 @@ mod tests {
assert!(body.contains("405 Not Allowed"), "got: {body}");
}
#[tokio::test]
async fn ws_upgrade_resolves_the_token_exactly_once() {
let provider = Arc::new(CountingProvider::new());
let app = build_router(
router_state(provider.clone() as Arc<dyn IdentityProvider>),
None,
);
let response = ws_upgrade_oneshot(app, "Bearer alk_test").await;
assert_eq!(
response.status(),
axum::http::StatusCode::SWITCHING_PROTOCOLS,
"a valid bearer token upgrades"
);
assert_eq!(
provider.resolutions(),
1,
"the WS upgrade path must resolve the token exactly once (SRV-11)"
);
let provider = Arc::new(CountingProvider::new());
let app = build_router(
router_state(provider.clone() as Arc<dyn IdentityProvider>),
None,
);
let response = ws_upgrade_oneshot(app, "Bearer alk_test").await;
drop(response);
assert_eq!(
provider.resolutions(),
1,
"exactly one resolution per upgrade request after any builder order"
);
}
#[tokio::test]
async fn ws_upgrade_without_token_is_rejected_401() {
let app = build_router(router_state(static_provider()), None);
let response = ws_upgrade_oneshot(app, "").await;
assert_eq!(
response.status(),
axum::http::StatusCode::UNAUTHORIZED,
"the WS route keeps its own enforced layer"
);
}
#[cfg(feature = "mcp")]
#[tokio::test]
async fn mcp_request_resolves_the_token_exactly_once() {
let provider = Arc::new(CountingProvider::new());
let app = build_router(
router_state(provider.clone() as Arc<dyn IdentityProvider>),
None,
);
let request = axum::http::Request::builder()
.method(axum::http::Method::POST)
.uri("/mcp")
.header(axum::http::header::HOST, "localhost")
.header(axum::http::header::AUTHORIZATION, "Bearer alk_test")
.header(axum::http::header::CONTENT_TYPE, "application/json")
.header(
axum::http::header::ACCEPT,
"application/json, text/event-stream",
)
.body(axum::body::Body::from(
serde_json::to_vec(&serde_json::json!({
"jsonrpc": "2.0", "id": 1, "method": "initialize",
"params": {
"protocolVersion": "2025-06-18",
"capabilities": {},
"clientInfo": { "name": "test-client", "version": "1.0.0" }
}
}))
.unwrap(),
))
.unwrap();
let response = tower::ServiceExt::oneshot(app, request).await.unwrap();
assert_eq!(
response.status(),
axum::http::StatusCode::OK,
"the /mcp initialize response comes back"
);
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap();
assert!(
String::from_utf8_lossy(&body).contains("alkhttp-to-mcp"),
"initialize response body, got: {}",
String::from_utf8_lossy(&body)
);
assert_eq!(
provider.resolutions(),
1,
"the /mcp path must resolve the token exactly once (SRV-11)"
);
}
#[tokio::test]
async fn method_mismatch_on_ws_upgrade_path_serves_decoy_405() {
let app = build_router(router_state(static_provider()), None);
let request = axum::http::Request::builder()
.method(axum::http::Method::POST)
.uri(WS_UPGRADE_PATH)
.body(axum::body::Body::empty())
.unwrap();
let response = get_with_bearer_with_method(app, request).await;
assert_eq!(
response.status(),
axum::http::StatusCode::METHOD_NOT_ALLOWED
);
let server = response
.headers()
.get(axum::http::header::SERVER)
.map(|v| v.to_str().unwrap().to_string());
assert_eq!(
server.as_deref(),
Some("nginx"),
"wrong-method probe on the WS path keeps the decoy shape after the reorder"
);
}
#[tokio::test]
async fn openapi_json_is_cached_and_generic_on_cache_miss() {
let adapter = HttpAdapter::new(static_provider(), empty_registry());