Move credential route proxy metadata to headers

This commit is contained in:
Winston Howes
2026-05-15 19:16:02 -07:00
parent 28ce16562d
commit dba0f036dc
2 changed files with 54 additions and 11 deletions
+32 -8
View File
@@ -33,6 +33,7 @@ use rama_http::Response;
use rama_http::StatusCode;
use rama_http::Uri;
use rama_http::header::HOST;
use rama_http::header::HeaderName;
use rama_http::layer::remove_header::RemoveRequestHeaderLayer;
use rama_http::layer::remove_header::RemoveResponseHeaderLayer;
use rama_http_backend::server::HttpServer;
@@ -85,6 +86,10 @@ enum MitmPolicyDecision {
const MITM_INSPECT_BODIES: bool = false;
const MITM_MAX_BODY_BYTES: usize = 4096;
const CREDENTIAL_ROUTE_CONNECTOR_ID_HEADER: &str =
"x-openai-internal-credential-route-connector-id";
const CREDENTIAL_ROUTE_LINK_ID_HEADER: &str = "x-openai-internal-credential-route-link-id";
const CREDENTIAL_ROUTE_REQUEST_URL_HEADER: &str = "x-openai-internal-credential-route-request-url";
impl std::fmt::Debug for MitmState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -234,7 +239,12 @@ async fn forward_request(req: Request, request_ctx: &MitmRequestContext) -> Resu
&& let Some(credentialed_route_proxy) = actions.credentialed_route_proxy.as_ref()
{
let (proxy_uri, proxy_authority) =
credentialed_route_proxy_request(credentialed_route_proxy, &original_request_uri)?;
credentialed_route_proxy_request(credentialed_route_proxy)?;
apply_credentialed_route_proxy_headers(
&mut parts.headers,
credentialed_route_proxy,
&original_request_uri,
)?;
parts.uri = proxy_uri;
parts
.headers
@@ -562,19 +572,33 @@ fn build_https_uri(authority: &str, path: &str) -> Result<Uri> {
fn credentialed_route_proxy_request(
action: &crate::mitm_hook::CredentialedRouteProxyAction,
request_uri: &Uri,
) -> Result<(Uri, String)> {
let mut proxy_url = Url::parse(&action.proxy_url)
let proxy_url = Url::parse(&action.proxy_url)
.with_context(|| format!("invalid credentialed route proxy URL {}", action.proxy_url))?;
proxy_url
.query_pairs_mut()
.append_pair("link_id", &action.link_id)
.append_pair("connector_id", &action.connector_id)
.append_pair("request_url", &request_uri.to_string());
let proxy_authority = url_authority_header_value(&proxy_url)?;
Ok((proxy_url.as_str().parse()?, proxy_authority))
}
fn apply_credentialed_route_proxy_headers(
headers: &mut HeaderMap,
action: &crate::mitm_hook::CredentialedRouteProxyAction,
request_uri: &Uri,
) -> Result<()> {
headers.insert(
HeaderName::from_static(CREDENTIAL_ROUTE_CONNECTOR_ID_HEADER),
HeaderValue::from_str(&action.connector_id)?,
);
headers.insert(
HeaderName::from_static(CREDENTIAL_ROUTE_LINK_ID_HEADER),
HeaderValue::from_str(&action.link_id)?,
);
headers.insert(
HeaderName::from_static(CREDENTIAL_ROUTE_REQUEST_URL_HEADER),
HeaderValue::from_str(&request_uri.to_string())?,
);
Ok(())
}
fn url_authority_header_value(url: &Url) -> Result<String> {
let host = url
.host_str()
+22 -3
View File
@@ -285,13 +285,32 @@ fn credentialed_route_proxy_request_rewrites_upstream_target() {
proxy_url: "http://localhost:8080/api/codex/credential_routes/proxy".to_string(),
};
let request_uri: Uri = "https://api.example.com/v1/items?limit=5".parse().unwrap();
let mut headers = HeaderMap::new();
headers.insert(
HeaderName::from_static(CREDENTIAL_ROUTE_REQUEST_URL_HEADER),
HeaderValue::from_static("https://spoofed.example.com"),
);
let (proxy_uri, proxy_authority) =
credentialed_route_proxy_request(&action, &request_uri).unwrap();
let (proxy_uri, proxy_authority) = credentialed_route_proxy_request(&action).unwrap();
apply_credentialed_route_proxy_headers(&mut headers, &action, &request_uri).unwrap();
assert_eq!(
proxy_uri.to_string(),
"http://localhost:8080/api/codex/credential_routes/proxy?link_id=link_123&connector_id=connector_123&request_url=https%3A%2F%2Fapi.example.com%2Fv1%2Fitems%3Flimit%3D5"
"http://localhost:8080/api/codex/credential_routes/proxy"
);
assert_eq!(proxy_authority, "localhost:8080");
assert_eq!(
headers.get(CREDENTIAL_ROUTE_CONNECTOR_ID_HEADER),
Some(&HeaderValue::from_static("connector_123"))
);
assert_eq!(
headers.get(CREDENTIAL_ROUTE_LINK_ID_HEADER),
Some(&HeaderValue::from_static("link_123"))
);
assert_eq!(
headers.get(CREDENTIAL_ROUTE_REQUEST_URL_HEADER),
Some(&HeaderValue::from_static(
"https://api.example.com/v1/items?limit=5"
))
);
}