Skip to main content

synapse_proxy/
proxy.rs

1//! Reverse-proxy pipeline: match (path_prefix + methods) → resolve context
2//! (gate require_context) → request transforms → forward (stream/hop-by-hop/cap)
3//! → response transforms → stream/replace.
4
5use std::sync::atomic::AtomicBool;
6use std::sync::Arc;
7
8use axum::body::Body;
9use axum::extract::{Request, State};
10use axum::http::{HeaderMap, StatusCode};
11use axum::response::{IntoResponse, Response};
12use axum::Json;
13
14use crate::builder::CompiledRoute;
15use crate::context::ContextStore;
16use crate::transform::{ProxyRequest, ProxyResponse, TransformError};
17
18const MAX_BODY_BYTES: usize = 64 * 1024 * 1024;
19
20pub const HOP_BY_HOP: &[&str] = &[
21    "connection",
22    "keep-alive",
23    "proxy-connection",
24    "transfer-encoding",
25    "upgrade",
26    "te",
27    "trailer",
28    "host",
29];
30
31pub fn is_hop_by_hop(name: &str) -> bool {
32    HOP_BY_HOP.contains(&name.to_ascii_lowercase().as_str())
33}
34
35#[derive(Clone)]
36pub struct AppState {
37    pub routes: Arc<Vec<CompiledRoute>>,
38    pub context: Arc<ContextStore>,
39    pub client: reqwest::Client,
40    pub shutting_down: Arc<AtomicBool>,
41    pub metrics: Arc<crate::metrics::Metrics>,
42}
43
44/// Longest `path_prefix` match, then narrow by `methods` (empty = any).
45pub fn match_route<'a>(
46    routes: &'a [CompiledRoute],
47    path: &str,
48    method: &str,
49) -> Option<&'a CompiledRoute> {
50    routes
51        .iter()
52        .filter(|r| path.starts_with(&r.path_prefix))
53        .filter(|r| r.methods.is_empty() || r.methods.iter().any(|m| m == method))
54        .max_by_key(|r| r.path_prefix.len())
55}
56
57fn err(status: StatusCode, error: &str, detail: String) -> Response {
58    (
59        status,
60        Json(serde_json::json!({ "error": error, "detail": detail })),
61    )
62        .into_response()
63}
64
65fn reject_response(e: TransformError) -> Response {
66    match e {
67        TransformError::Reject {
68            status,
69            error,
70            detail,
71        } => err(status, &error, detail),
72        TransformError::Internal(m) => err(StatusCode::INTERNAL_SERVER_ERROR, "transform_error", m),
73    }
74}
75
76fn reject_status(e: &TransformError) -> u16 {
77    match e {
78        TransformError::Reject { status, .. } => status.as_u16(),
79        TransformError::Internal(_) => 500,
80    }
81}
82
83pub async fn handler(State(state): State<AppState>, req: Request) -> Response {
84    let started = std::time::Instant::now();
85    let (parts, body) = req.into_parts();
86    let path = parts.uri.path().to_string();
87    let method = parts.method.as_str().to_string();
88
89    let Some(route) = match_route(&state.routes, &path, &method) else {
90        let secs = started.elapsed().as_secs_f64();
91        state.metrics.record("none", &method, 404, "no_route", secs);
92        return err(
93            StatusCode::NOT_FOUND,
94            "no_route",
95            format!("no route matches '{method} {path}'"),
96        );
97    };
98
99    let route_label = route.name.clone();
100
101    let ctx = state.context.resolve();
102    for key in &route.require_context {
103        if !ctx.contains(key) {
104            let secs = started.elapsed().as_secs_f64();
105            state
106                .metrics
107                .record(&route_label, &method, 503, "context_unbound", secs);
108            return err(
109                StatusCode::SERVICE_UNAVAILABLE,
110                "request_failed",
111                "context not bound".into(),
112            );
113        }
114    }
115
116    let bytes = match axum::body::to_bytes(body, MAX_BODY_BYTES).await {
117        Ok(b) => b.to_vec(),
118        Err(e) => {
119            return err(
120                StatusCode::PAYLOAD_TOO_LARGE,
121                "body_too_large",
122                e.to_string(),
123            )
124        }
125    };
126
127    // Inbound headers minus hop-by-hop become the transform-visible header set.
128    let mut headers = HeaderMap::new();
129    for (n, v) in parts.headers.iter() {
130        if !is_hop_by_hop(n.as_str()) {
131            headers.insert(n.clone(), v.clone());
132        }
133    }
134    let query = parts.uri.query().map(str::to_string);
135    let mut preq = ProxyRequest::from_parts(
136        parts.method.clone(),
137        path.clone(),
138        query.clone(),
139        headers,
140        bytes,
141    );
142
143    for t in &route.request {
144        if let Err(e) = t.apply(&ctx, &mut preq).await {
145            let secs = started.elapsed().as_secs_f64();
146            state.metrics.record(
147                &route_label,
148                &method,
149                reject_status(&e),
150                "transform_rejected",
151                secs,
152            );
153            state.metrics.transform_error(&route_label, "request");
154            return reject_response(e);
155        }
156    }
157
158    // Build the upstream URL (prefix strip + query).
159    let rest = if route.strip_prefix {
160        preq.path
161            .strip_prefix(&route.path_prefix)
162            .unwrap_or(&preq.path)
163    } else {
164        &preq.path
165    };
166    let base = route.upstream.trim_end_matches('/');
167    let mut url = format!("{base}{rest}");
168    if let Some(q) = preq.query.as_deref().filter(|q| !q.is_empty()) {
169        url.push('?');
170        url.push_str(q);
171    }
172
173    let method_val = preq.method.clone();
174    let out_headers = preq.headers.clone();
175    let out_body = preq.into_body_bytes();
176
177    let upstream = state
178        .client
179        .request(method_val, &url)
180        .headers(out_headers)
181        .body(out_body)
182        .send()
183        .await;
184    let resp = match upstream {
185        Ok(r) => r,
186        Err(e) => {
187            let secs = started.elapsed().as_secs_f64();
188            state
189                .metrics
190                .record(&route_label, &method, 502, "upstream_error", secs);
191            state.metrics.upstream_error(&route_label, "send");
192            return err(StatusCode::BAD_GATEWAY, "request_failed", e.to_string());
193        }
194    };
195
196    // Response transforms (status/headers/replacement only in v1).
197    let status = StatusCode::from_u16(resp.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
198    let mut resp_headers = HeaderMap::new();
199    for (n, v) in resp.headers().iter() {
200        if !is_hop_by_hop(n.as_str()) {
201            resp_headers.insert(n.clone(), v.clone());
202        }
203    }
204    let mut presp = ProxyResponse::new(status, resp_headers);
205    for t in &route.response {
206        if let Err(e) = t.apply(&ctx, &mut presp).await {
207            let secs = started.elapsed().as_secs_f64();
208            state.metrics.record(
209                &route_label,
210                &method,
211                reject_status(&e),
212                "transform_rejected",
213                secs,
214            );
215            state.metrics.transform_error(&route_label, "response");
216            return reject_response(e);
217        }
218    }
219
220    let final_status = presp.status.as_u16();
221    let secs = started.elapsed().as_secs_f64();
222    state
223        .metrics
224        .record(&route_label, &method, final_status, "forwarded", secs);
225
226    if let Some(replacement) = presp.replacement() {
227        return (presp.status, Json(replacement.clone())).into_response();
228    }
229
230    // No body replacement → stream the upstream body unchanged.
231    let mut builder = Response::builder().status(presp.status);
232    for (n, v) in presp.headers.iter() {
233        builder = builder.header(n, v);
234    }
235    builder
236        .body(Body::from_stream(resp.bytes_stream()))
237        .unwrap_or_else(|e| err(StatusCode::BAD_GATEWAY, "request_failed", e.to_string()))
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243    use crate::builder::CompiledRoute;
244
245    fn route(prefix: &str, methods: &[&str]) -> CompiledRoute {
246        CompiledRoute {
247            name: prefix.into(),
248            path_prefix: prefix.into(),
249            upstream: "http://u".into(),
250            strip_prefix: false,
251            methods: methods.iter().map(|s| s.to_string()).collect(),
252            require_context: vec![],
253            request: vec![],
254            response: vec![],
255        }
256    }
257
258    #[test]
259    fn matches_longest_prefix_and_method() {
260        let routes = vec![route("/v1", &[]), route("/v1/llm", &["POST"])];
261        assert_eq!(
262            match_route(&routes, "/v1/llm/x", "POST")
263                .unwrap()
264                .path_prefix,
265            "/v1/llm"
266        );
267        // method mismatch on the longer route → falls back to the any-method route
268        assert_eq!(
269            match_route(&routes, "/v1/llm/x", "GET")
270                .unwrap()
271                .path_prefix,
272            "/v1"
273        );
274        assert!(match_route(&routes, "/nope", "GET").is_none());
275    }
276}