diff --git a/crates/aisix-proxy/src/error.rs b/crates/aisix-proxy/src/error.rs index 0033e18a..fa443ad3 100644 --- a/crates/aisix-proxy/src/error.rs +++ b/crates/aisix-proxy/src/error.rs @@ -208,6 +208,15 @@ pub enum ProxyError { ModelIpRestricted(String), #[error("request payload is invalid: {0}")] InvalidRequest(String), + /// A non-WebSocket request reached the WebSocket-only realtime + /// endpoint. Carries the upgrade layer's own classification — 400 for + /// malformed upgrade headers, 426 for a connection that cannot + /// upgrade, 405 for a HEAD request (axum's `get()` also serves HEAD, + /// so the extractor's method check is reachable) — plus its + /// per-variant reason, so the refusal keeps both the status and the + /// diagnostic the bare rejection used to carry. + #[error("this endpoint requires a WebSocket upgrade: {detail}")] + WebSocketUpgradeRequired { status: StatusCode, detail: String }, #[error("no bridge registered for provider")] ProviderUnavailable, /// Every routing candidate was excluded by the runtime status layer @@ -287,6 +296,7 @@ impl ProxyError { ProxyError::ModelNotFound(_) => StatusCode::NOT_FOUND, ProxyError::VideoNotFound(_) => StatusCode::NOT_FOUND, ProxyError::InvalidRequest(_) => StatusCode::BAD_REQUEST, + ProxyError::WebSocketUpgradeRequired { status, .. } => *status, ProxyError::ProviderUnavailable => StatusCode::SERVICE_UNAVAILABLE, ProxyError::AllCandidatesUnavailable { .. } => StatusCode::SERVICE_UNAVAILABLE, ProxyError::ContentFiltered(_) => StatusCode::UNPROCESSABLE_ENTITY, @@ -317,6 +327,7 @@ impl ProxyError { ProxyError::ModelNotFound(_) => "model_not_found", ProxyError::VideoNotFound(_) => "video_not_found", ProxyError::InvalidRequest(_) => "invalid_request_error", + ProxyError::WebSocketUpgradeRequired { .. } => "websocket_upgrade_required", ProxyError::RequestTooLarge { .. } => "invalid_request_error", ProxyError::ProviderUnavailable => "provider_unavailable", ProxyError::AllCandidatesUnavailable { .. } => "all_candidates_unavailable", @@ -478,6 +489,7 @@ impl IntoResponse for ProxyError { fn into_response(self) -> Response { let status = self.status(); let retry_after = self.retry_after_secs(); + let upgrade_reject = matches!(self, ProxyError::WebSocketUpgradeRequired { .. }); let body = self.envelope(); let mut response = (status, Json(body)).into_response(); if let Some(secs) = retry_after { @@ -485,6 +497,25 @@ impl IntoResponse for ProxyError { response.headers_mut().insert("retry-after", value); } } + // RFC 9110: a 426 must name the protocol to switch to (§15.5.22) + // and a 405 must list the allowed methods (§15.5.6; GET implies + // HEAD on this route). axum's bare rejection omitted both, but the + // response is the gateway's own now. + if upgrade_reject { + match status { + StatusCode::UPGRADE_REQUIRED => { + response + .headers_mut() + .insert("upgrade", HeaderValue::from_static("websocket")); + } + StatusCode::METHOD_NOT_ALLOWED => { + response + .headers_mut() + .insert("allow", HeaderValue::from_static("GET, HEAD")); + } + _ => {} + } + } response } } diff --git a/crates/aisix-proxy/src/realtime.rs b/crates/aisix-proxy/src/realtime.rs index 1c20953a..9538e45d 100644 --- a/crates/aisix-proxy/src/realtime.rs +++ b/crates/aisix-proxy/src/realtime.rs @@ -91,16 +91,34 @@ async fn connect_upstream( pub(crate) async fn realtime( State(state): State, + method: Method, Query(params): Query>, headers: HeaderMap, client: ClientContext, - ws: WebSocketUpgrade, + ws: Result, ) -> Response { let request_id = client.request_id.clone(); let started = Instant::now(); - match prepare(&state, ¶ms, &headers, &client).await { - Ok(prep) => { + // A non-WebSocket request (plain GET, malformed upgrade headers, a + // connection that cannot upgrade, or a HEAD — `get()` serves HEAD + // too) used to get axum's bare rejection — no access log, no metrics, + // no usage event, no envelope; the same silent class #863/#880/#884 + // collected (#885). Map it into this endpoint's normal error arm, + // keeping axum's own status classification (400 / 426 / 405) and its + // per-variant diagnostic. + let outcome = match ws { + Ok(ws) => prepare(&state, ¶ms, &headers, &client) + .await + .map(|prep| (ws, prep)), + Err(rejection) => Err(crate::error::ProxyError::WebSocketUpgradeRequired { + status: rejection.status(), + detail: rejection.body_text(), + }), + }; + + match outcome { + Ok((ws, prep)) => { let state2 = state.clone(); let client2 = client.clone(); // `on_upgrade` runs the session on a detached task, so the @@ -119,13 +137,24 @@ pub(crate) async fn realtime( Err(err) => { let status = err.status().as_u16(); emit_access_log( - &Method::GET, + &method, status, started.elapsed(), &request_id, None, Some(&err), ); + // Count the refusal like every other pre-dispatch rejection + // (unresolved labels, same as `reject_before_dispatch`) — logs + // and the request-rate metrics must not disagree about whether + // these requests exist. + state.metrics.record_request( + "unknown", + crate::usage_attr::UNRESOLVED_MODEL_LABEL, + status, + RequestOutcome::from_status(status), + started.elapsed(), + ); crate::usage_attr::emit_error_usage_event( &state, "realtime", @@ -1080,4 +1109,136 @@ mod tests { .expect_err("handshake must fail on model ACL"); assert!(err.to_string().contains("403"), "got: {err}"); } + + /// State + router + usage receiver for driving the endpoint's + /// REJECTION paths with `oneshot` (no live connection needed — the + /// point is that no upgrade happens). + fn oneshot_router() -> ( + axum::Router, + crate::ProxyState, + tokio::sync::mpsc::Receiver, + ) { + let (tx, rx) = tokio::sync::mpsc::channel::(4); + let state = crate::ProxyState::new( + SnapshotHandle::new(AisixSnapshot::new()), + Arc::new(Hub::new()), + &cfg(), + ) + .without_cache() + .with_usage_sink(UsageSink::new(tx)); + (crate::build_router(state.clone()), state, rx) + } + + async fn body_json(response: axum::response::Response) -> serde_json::Value { + let bytes = axum::body::to_bytes(response.into_body(), 1 << 16) + .await + .unwrap_or_default(); + serde_json::from_slice(&bytes).expect("an error envelope, not a bare rejection body") + } + + #[tokio::test] + async fn non_websocket_request_is_recorded_and_enveloped() { + use tower::ServiceExt as _; + // Pre-#885 a plain GET (no upgrade headers) got axum's bare + // rejection: nothing in the access log, metrics, or the usage + // pipeline. It now takes this endpoint's normal error arm — + // envelope + usage event + request metrics — keeping axum's 400 + // classification for bad/missing upgrade headers. + let (router, state, mut rx) = oneshot_router(); + let response = router + .oneshot( + axum::http::Request::get("/v1/realtime?model=probe-model") + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .expect("router responds"); + assert_eq!(response.status(), axum::http::StatusCode::BAD_REQUEST); + let json = body_json(response).await; + assert_eq!(json["error"]["type"], "websocket_upgrade_required"); + + let event = rx.try_recv().expect("the refusal is recorded"); + assert_eq!(event.status_code, 400); + assert_eq!(event.inbound_protocol, "realtime"); + // Auth runs inside prepare(), which a rejected upgrade never + // reaches — no key is attributed; the requested model rides along. + assert_eq!(event.api_key_id, ""); + assert_eq!(event.requested_model, "probe-model"); + + // Logs and the request-rate metrics must not disagree about + // whether these requests exist. + let scrape = state.metrics.render(); + assert!( + scrape.contains(r#"status="400""#) && scrape.contains(r#"model="unresolved""#), + "the refusal must be counted, got: {scrape}" + ); + } + + #[tokio::test] + async fn non_upgradable_connection_keeps_its_426() { + use tower::ServiceExt as _; + // Correct WebSocket headers over a connection that cannot upgrade + // (a `oneshot` request carries no hyper upgrade extension) is + // axum's ConnectionNotUpgradable — 426 Upgrade Required. The + // status must survive the envelope mapping rather than being + // flattened to 400, and the response must name the protocol to + // switch to (RFC 9110 §15.5.22). + let (router, _state, mut rx) = oneshot_router(); + let response = router + .oneshot( + axum::http::Request::get("/v1/realtime") + .header("connection", "upgrade") + .header("upgrade", "websocket") + .header("sec-websocket-version", "13") + .header("sec-websocket-key", "dGhlIHNhbXBsZSBub25jZQ==") + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .expect("router responds"); + assert_eq!( + response.status(), + axum::http::StatusCode::UPGRADE_REQUIRED, + "axum's 426 classification must survive" + ); + assert_eq!( + response + .headers() + .get("upgrade") + .and_then(|v| v.to_str().ok()), + Some("websocket") + ); + let json = body_json(response).await; + assert_eq!(json["error"]["type"], "websocket_upgrade_required"); + assert_eq!(rx.try_recv().expect("recorded").status_code, 426); + } + + #[tokio::test] + async fn head_request_keeps_its_405_and_allow_header() { + use tower::ServiceExt as _; + // axum's `get()` also serves HEAD, so a HEAD request reaches the + // extractor's method check — 405, with the Allow header RFC 9110 + // §15.5.6 requires, and recorded like every other refusal. + let (router, _state, mut rx) = oneshot_router(); + let response = router + .oneshot( + axum::http::Request::head("/v1/realtime") + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .expect("router responds"); + assert_eq!( + response.status(), + axum::http::StatusCode::METHOD_NOT_ALLOWED + ); + assert_eq!( + response + .headers() + .get("allow") + .and_then(|v| v.to_str().ok()), + Some("GET, HEAD") + ); + assert_eq!(rx.try_recv().expect("recorded").status_code, 405); + } } diff --git a/tests/e2e/src/cases/realtime-ws-e2e.test.ts b/tests/e2e/src/cases/realtime-ws-e2e.test.ts index 4a5cdcfd..c7fa3714 100644 --- a/tests/e2e/src/cases/realtime-ws-e2e.test.ts +++ b/tests/e2e/src/cases/realtime-ws-e2e.test.ts @@ -182,6 +182,19 @@ describe("realtime e2e: /v1/realtime WebSocket relay (#721)", () => { ws.close(); }); + test("a plain http GET answers the envelope, not a bare rejection", async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + // No upgrade headers: pre-#885 this got axum's bare 400 with no + // telemetry. It must now wear the endpoint's error envelope. + const res = await fetch(`${app.proxyUrl}/v1/realtime`); + expect(res.status).toBe(400); + const body = (await res.json()) as { error?: { type?: string } }; + expect(body.error?.type).toBe("websocket_upgrade_required"); + }); + test("bad credentials reject the upgrade handshake", async (ctx) => { if (!etcdReachable || !app) { ctx.skip();