From f63520c961b3a14523787f48389abe85ef55d36d Mon Sep 17 00:00:00 2001 From: AprilNEA Date: Thu, 23 Apr 2026 20:55:00 +0800 Subject: [PATCH 1/2] fix(proxy): make ampcode passthrough fully transparent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The /api/* and /v0/management/* catch-all proxy was rewriting the amp CLI's own headers (stripping Authorization/X-Api-Key, injecting byokey's amp token) and buffering the entire upstream response before returning it. Both broke the amp `read` tool: the rewritten auth caused thread operations to run as the wrong principal, and buffering prevented the SSE turn-continuation from streaming back to the client. Forward the amp CLI's request verbatim (only hop-by-hop + host stripped per RFC 7230) and stream the response body back. /api/provider/* is unaffected — byokey still owns AI traffic. The forward_headers middleware is deleted since no other route needed it. --- crates/proxy/src/handler/amp/provider.rs | 78 ++++++++++++------- crates/proxy/src/middleware/forward.rs | 97 ------------------------ crates/proxy/src/middleware/mod.rs | 1 - crates/proxy/src/router.rs | 15 ++-- 4 files changed, 55 insertions(+), 136 deletions(-) delete mode 100644 crates/proxy/src/middleware/forward.rs diff --git a/crates/proxy/src/handler/amp/provider.rs b/crates/proxy/src/handler/amp/provider.rs index f935eb1..5d2e931 100644 --- a/crates/proxy/src/handler/amp/provider.rs +++ b/crates/proxy/src/handler/amp/provider.rs @@ -24,7 +24,6 @@ use futures_util::TryStreamExt as _; use serde_json::Value; use std::{collections::HashMap, sync::Arc}; -use crate::middleware::forward::ForwardedHeaders; use crate::util::stream::{CodexParser, GeminiParser, response_to_stream, tap_usage_stream}; use crate::util::{bad_gateway, extract_usage, sse_response, upstream_error}; use crate::{AppState, error::ApiError}; @@ -533,12 +532,34 @@ fn byte_stream_to_gemini_sse( }) } +/// Hop-by-hop headers per RFC 7230 §6.1 — must not be forwarded by proxies. +const HOP_BY_HOP: &[&str] = &[ + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailers", + "transfer-encoding", + "upgrade", +]; + /// Transparent proxy to `ampcode.com` — used for both `/api/{*path}` and -/// `/v0/management/{*path}`. Takes the original URI path directly so a -/// single handler covers all non-provider amp routes. +/// `/v0/management/{*path}`. +/// +/// Forwards the request verbatim: the client's headers (including +/// `Authorization`, cookies, user-agent, etc.) and body are passed through +/// untouched apart from the HTTP hop-by-hop headers and `host`, which any +/// proxy must drop for correctness. The response body is streamed back +/// rather than buffered so SSE and long-poll endpoints (e.g. thread +/// continuation after a tool result) work end-to-end. +/// +/// byokey deliberately does NOT inject its own amp token here: the amp CLI +/// manages its own user session, and these non-AI endpoints +/// (threads, telemetry, auth, management) must run as the logged-in user. pub async fn ampcode_proxy( State(state): State>, - axum::extract::Extension(fwd): axum::extract::Extension, + headers: axum::http::HeaderMap, method: Method, uri: axum::http::Uri, body: Bytes, @@ -549,22 +570,24 @@ pub async fn ampcode_proxy( None => format!("{AMP_BACKEND}{path}"), }; - let debug = path.ends_with("/internal") && tracing::enabled!(tracing::Level::DEBUG); - if debug { - let req_body = std::str::from_utf8(&body) - .ok() - .and_then(|s| serde_json::from_str::(s).ok()) - .map_or_else( - || format!("{body:?}"), - |v| serde_json::to_string_pretty(&v).unwrap_or_default(), - ); - tracing::debug!(%method, %url, body = %req_body, "ampcode proxy request"); + let mut fwd = rquest::header::HeaderMap::new(); + for (name, value) in &headers { + let name_str = name.as_str(); + if HOP_BY_HOP.contains(&name_str) || name_str == "host" { + continue; + } + if let (Ok(n), Ok(v)) = ( + rquest::header::HeaderName::from_bytes(name.as_ref()), + rquest::header::HeaderValue::from_bytes(value.as_bytes()), + ) { + fwd.insert(n, v); + } } let resp = match state .http .request(method, url) - .headers(fwd.headers) + .headers(fwd) .body(body) .send() .await @@ -577,6 +600,10 @@ pub async fn ampcode_proxy( let mut resp_headers = axum::http::HeaderMap::new(); for (name, value) in resp.headers() { + let name_str = name.as_str(); + if HOP_BY_HOP.contains(&name_str) { + continue; + } if let (Ok(n), Ok(v)) = ( axum::http::HeaderName::from_bytes(name.as_ref()), axum::http::HeaderValue::from_bytes(value.as_bytes()), @@ -585,20 +612,15 @@ pub async fn ampcode_proxy( } } - let body_bytes = resp.bytes().await.unwrap_or_default(); - - if debug { - let resp_body = std::str::from_utf8(&body_bytes) - .ok() - .and_then(|s| serde_json::from_str::(s).ok()) - .map_or_else( - || format!("{body_bytes:?}"), - |v| serde_json::to_string_pretty(&v).unwrap_or_default(), - ); - tracing::debug!(%status, body = %resp_body, "ampcode proxy response"); - } + let stream = resp + .bytes_stream() + .map_err(|e| std::io::Error::other(e.to_string())); + let body = axum::body::Body::from_stream(stream); - (status, resp_headers, body_bytes).into_response() + let mut response = Response::new(body); + *response.status_mut() = status; + *response.headers_mut() = resp_headers; + response } #[cfg(test)] diff --git a/crates/proxy/src/middleware/forward.rs b/crates/proxy/src/middleware/forward.rs deleted file mode 100644 index 714d10f..0000000 --- a/crates/proxy/src/middleware/forward.rs +++ /dev/null @@ -1,97 +0,0 @@ -//! Header-forwarding middleware for amp proxy routes. -//! -//! Resolves amp auth, filters hop-by-hop / fingerprint headers, injects -//! `Authorization` + `X-Api-Key`, and stores the result as [`ForwardedHeaders`] -//! in request extensions so handlers can use it directly. - -use axum::{extract::Request, middleware::Next, response::Response}; -use byokey_types::ProviderId; -use std::sync::Arc; - -use crate::AppState; - -const HOP_BY_HOP: &[&str] = &[ - "connection", - "keep-alive", - "proxy-authenticate", - "proxy-authorization", - "te", - "trailers", - "transfer-encoding", - "upgrade", -]; - -const CLIENT_AUTH_HEADERS: &[&str] = &["authorization", "x-api-key", "x-goog-api-key"]; - -const FINGERPRINT_HEADERS: &[&str] = &[ - "x-forwarded-for", - "x-forwarded-host", - "x-forwarded-proto", - "x-real-ip", - "forwarded", - "via", - "priority", -]; - -/// Prepared upstream headers with amp auth already injected. -/// Stored in request extensions by [`forward_headers_middleware`]. -#[derive(Clone)] -pub struct ForwardedHeaders { - pub headers: rquest::header::HeaderMap, -} - -pub async fn forward_headers_middleware( - axum::extract::State(state): axum::extract::State>, - mut request: Request, - next: Next, -) -> Response { - let config = state.config.load(); - let amp_token = state.auth.get_token(&ProviderId::Amp).await.ok(); - let strip_auth = amp_token.is_some() || config.amp.upstream_key.is_some(); - - let mut out = rquest::header::HeaderMap::new(); - for (name, value) in request.headers() { - let name_str = name.as_str(); - if HOP_BY_HOP.contains(&name_str) || name_str == "host" { - continue; - } - if strip_auth && CLIENT_AUTH_HEADERS.contains(&name_str) { - continue; - } - if FINGERPRINT_HEADERS.contains(&name_str) - || name_str.starts_with("sec-ch-ua-") - || name_str.starts_with("sec-fetch-") - { - continue; - } - if let (Ok(n), Ok(v)) = ( - rquest::header::HeaderName::from_bytes(name.as_ref()), - rquest::header::HeaderValue::from_bytes(value.as_bytes()), - ) { - out.insert(n, v); - } - } - - if let Some(token) = &_token { - inject_amp_auth(&mut out, &token.access_token); - } else if let Some(key) = &config.amp.upstream_key { - inject_amp_auth(&mut out, key); - } - - request - .extensions_mut() - .insert(ForwardedHeaders { headers: out }); - next.run(request).await -} - -fn inject_amp_auth(headers: &mut rquest::header::HeaderMap, token: &str) { - if let (Ok(n_auth), Ok(v_auth), Ok(n_apikey), Ok(v_apikey)) = ( - rquest::header::HeaderName::from_bytes(b"authorization"), - rquest::header::HeaderValue::from_str(&format!("Bearer {token}")), - rquest::header::HeaderName::from_bytes(b"x-api-key"), - rquest::header::HeaderValue::from_str(token), - ) { - headers.insert(n_auth, v_auth); - headers.insert(n_apikey, v_apikey); - } -} diff --git a/crates/proxy/src/middleware/mod.rs b/crates/proxy/src/middleware/mod.rs index 5f7296b..6fc464f 100644 --- a/crates/proxy/src/middleware/mod.rs +++ b/crates/proxy/src/middleware/mod.rs @@ -1,4 +1,3 @@ //! Axum middleware layers for the proxy. pub mod dump; -pub mod forward; diff --git a/crates/proxy/src/router.rs b/crates/proxy/src/router.rs index cd0e37a..1a58aac 100644 --- a/crates/proxy/src/router.rs +++ b/crates/proxy/src/router.rs @@ -96,12 +96,11 @@ fn common_layers(router: Router) -> Router { /// `/byokey.amp.AmpService/{Method}` — local byokey management over /// `ConnectRPC` (fallback service). /// -/// The amp routes are wrapped in [`forward_headers_middleware`] to strip -/// client auth and inject the amp upstream token. The middleware is -/// scoped to that sub-router only via `.layer()` before `.merge()`, so -/// REST and `ConnectRPC` routes are unaffected. +/// `/api/provider/*` is handled locally (byokey's own AI providers); +/// every other `/api/*` and `/v0/management/*` request is a transparent +/// passthrough to `ampcode.com` with the amp CLI's own headers — byokey +/// only touches traffic it generates itself. pub fn make_router(state: Arc) -> Router { - // Amp-specific routes with forward_headers_middleware scoped to them. let amp_routes = Router::new() .route("/auth/cli-login", get(amp::cli_login_redirect)) .route("/v1/login", get(amp::login_redirect)) @@ -122,11 +121,7 @@ pub fn make_router(state: Arc) -> Router { "/api/provider/google/v1beta/models/{action}", post(amp::provider::gemini_native_passthrough), ) - .route("/api/{*path}", any(amp::provider::ampcode_proxy)) - .layer(middleware::from_fn_with_state( - state.clone(), - crate::middleware::forward::forward_headers_middleware, - )); + .route("/api/{*path}", any(amp::provider::ampcode_proxy)); // REST AI proxy routes. let rest_routes = Router::new() From d938d64ebaadbd13fbbaf52ecc02a3e3ef0ed4d2 Mon Sep 17 00:00:00 2001 From: AprilNEA Date: Thu, 23 Apr 2026 21:03:20 +0800 Subject: [PATCH 2/2] fix(proxy): strip content-encoding/length when forwarding amp responses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit rquest is built with `features = ["full"]` so it auto-decompresses gzip/brotli/zstd response bodies before `bytes_stream()` yields them. The prior passthrough still forwarded upstream's `Content-Encoding` and `Content-Length` headers, which described the *compressed* payload — so the amp CLI saw `Content-Encoding: gzip` and tried to decode already-decoded bytes, got garbage, silently failed, and fell back to a degraded mode policy where `read`/`bash` aren't allowed in smart mode. Strip both headers on the catch-all response path so axum sends the decompressed body as chunked without a misleading encoding hint. --- crates/proxy/src/handler/amp/provider.rs | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/crates/proxy/src/handler/amp/provider.rs b/crates/proxy/src/handler/amp/provider.rs index 5d2e931..5c1317e 100644 --- a/crates/proxy/src/handler/amp/provider.rs +++ b/crates/proxy/src/handler/amp/provider.rs @@ -598,10 +598,18 @@ pub async fn ampcode_proxy( let status = StatusCode::from_u16(resp.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY); + // rquest auto-decompresses gzip/brotli/zstd. The upstream's + // `content-encoding` and `content-length` describe the *compressed* + // payload, so forwarding them alongside the already-decompressed body + // would make the amp CLI try to decode twice (garbage → silent parse + // failure → degraded mode policy, losing read/bash in smart mode). let mut resp_headers = axum::http::HeaderMap::new(); for (name, value) in resp.headers() { let name_str = name.as_str(); - if HOP_BY_HOP.contains(&name_str) { + if HOP_BY_HOP.contains(&name_str) + || name_str == "content-encoding" + || name_str == "content-length" + { continue; } if let (Ok(n), Ok(v)) = (