1use 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
44pub 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 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 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 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 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 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}