Skip to main content

phrona_api/
lib.rs

1pub mod frontend;
2pub mod grounding;
3pub mod metrics;
4pub mod tavily;
5pub mod tools;
6
7use std::collections::HashMap;
8use std::net::{IpAddr, SocketAddr};
9use std::sync::Arc;
10use std::time::{Duration, Instant};
11
12use axum::extract::rejection::{BytesRejection, FailedToBufferBody, JsonRejection, QueryRejection};
13use axum::extract::{
14    ConnectInfo, DefaultBodyLimit, FromRequest, FromRequestParts, Query, Request, State,
15};
16use axum::http::{HeaderMap, StatusCode};
17use axum::middleware::{self, Next};
18use axum::response::{IntoResponse, Response};
19use axum::routing::{get, post};
20use axum::{Json, Router};
21use parking_lot::Mutex;
22use serde::Deserialize;
23use serde::de::DeserializeOwned;
24use serde_json::{Value, json};
25use tower_http::cors::CorsLayer;
26use tower_http::trace::TraceLayer;
27
28use phrona::engine;
29use phrona::error::{ErrorKind as PhronaErrorKind, ErrorScope};
30use phrona::models::{Category, SearchResponse, TimeRange};
31use phrona::{PhronaConfig, SearchClient, SearchOptions, suggest, suggest_all};
32
33/// Fixed-window rate-limit bucket for one client.
34struct RateWindow {
35    started: Instant,
36    count: u32,
37}
38
39pub struct AppState {
40    pub client: SearchClient,
41    pub started: Instant,
42    pub api_key: Option<String>,
43    /// Upper bound applied to `max_results` (from `search.max_results_limit`).
44    pub max_results_limit: usize,
45    /// Requests allowed per window per client; `0` disables the limiter.
46    pub rate_limit_per_minute: u32,
47    /// Maximum accepted request body size in bytes.
48    pub max_body_bytes: u64,
49    /// Reverse-proxy IPs whose `X-Forwarded-For` header is trusted for
50    /// client IP extraction (see [`client_ip`]).
51    pub trusted_proxies: Vec<IpAddr>,
52    rate: Mutex<HashMap<Option<IpAddr>, RateWindow>>,
53}
54
55impl AppState {
56    pub fn new(
57        client: SearchClient,
58        api_key: Option<String>,
59        max_results_limit: usize,
60        rate_limit_per_minute: u32,
61        max_body_bytes: u64,
62        trusted_proxies: Vec<IpAddr>,
63    ) -> Self {
64        Self {
65            client,
66            started: Instant::now(),
67            api_key,
68            max_results_limit,
69            rate_limit_per_minute,
70            max_body_bytes,
71            trusted_proxies,
72            rate: Mutex::new(HashMap::new()),
73        }
74    }
75
76    /// Constant-time API-key comparison: a configured key is only ever
77    /// compared with XOR folds of equal-length byte buffers, never with
78    /// early-exit string equality, so timing cannot leak key prefixes.
79    pub fn authorized(&self, key: Option<&str>) -> bool {
80        match (&self.api_key, key) {
81            (None, _) => true,
82            (Some(want), Some(got)) => {
83                let w = want.as_bytes();
84                let g = got.as_bytes();
85                w.len() == g.len() && w.iter().zip(g).fold(0u8, |acc, (a, b)| acc | (a ^ b)) == 0
86            }
87            (Some(_), None) => false,
88        }
89    }
90
91    /// Fixed-window rate limit keyed on the client IP (falling back to a
92    /// single global bucket when the client address is unknown, e.g. in
93    /// unit tests). Returns `true` when the request is allowed.
94    pub fn check_rate(&self, ip: Option<IpAddr>) -> bool {
95        if self.rate_limit_per_minute == 0 {
96            return true;
97        }
98        let mut rate = self.rate.lock();
99        // Bounded memory: under public traffic (or IPv6 address rotation) the
100        // per-IP bucket map would otherwise grow forever. Once it exceeds a
101        // threshold, sweep windows whose 60s window has already elapsed.
102        const RATE_BUCKET_LIMIT: usize = 10_000;
103        if rate.len() > RATE_BUCKET_LIMIT {
104            rate.retain(|_, w| w.started.elapsed() < Duration::from_secs(60));
105        }
106        let window = rate.entry(ip).or_insert(RateWindow {
107            started: Instant::now(),
108            count: 0,
109        });
110        if window.started.elapsed() >= Duration::from_secs(60) {
111            window.started = Instant::now();
112            window.count = 0;
113        }
114        if window.count >= self.rate_limit_per_minute {
115            return false;
116        }
117        window.count += 1;
118        true
119    }
120}
121
122pub struct AppError(ErrorKind);
123
124enum ErrorKind {
125    BadRequest(String),
126    Unauthorized,
127    RateLimited(String),
128    BodyTooLarge(u64),
129    Internal(phrona::Error),
130}
131
132impl AppError {
133    fn bad_request(msg: impl Into<String>) -> Self {
134        Self(ErrorKind::BadRequest(msg.into()))
135    }
136
137    fn unauthorized() -> Self {
138        Self(ErrorKind::Unauthorized)
139    }
140
141    fn rate_limited(limit: u32) -> Self {
142        Self(ErrorKind::RateLimited(format!(
143            "rate limit exceeded: at most {limit} requests per minute"
144        )))
145    }
146
147    fn body_too_large(max: u64) -> Self {
148        Self(ErrorKind::BodyTooLarge(max))
149    }
150}
151
152impl IntoResponse for AppError {
153    fn into_response(self) -> Response {
154        let (status, body) = match self.0 {
155            ErrorKind::BadRequest(msg) => (StatusCode::BAD_REQUEST, json!({"error": msg})),
156            ErrorKind::Unauthorized => (
157                StatusCode::UNAUTHORIZED,
158                json!({"error": "invalid api key"}),
159            ),
160            ErrorKind::RateLimited(msg) => (StatusCode::TOO_MANY_REQUESTS, json!({"error": msg})),
161            ErrorKind::BodyTooLarge(max) => (
162                StatusCode::PAYLOAD_TOO_LARGE,
163                json!({"error": format!("request body exceeds the {max}-byte limit")}),
164            ),
165            ErrorKind::Internal(e) => {
166                tracing::error!("search failed: {e}");
167                let status = if matches!(e.kind(), PhronaErrorKind::RateLimited { .. }) {
168                    StatusCode::TOO_MANY_REQUESTS
169                } else {
170                    match e.scope() {
171                        ErrorScope::Query => StatusCode::BAD_REQUEST,
172                        ErrorScope::Internal => StatusCode::INTERNAL_SERVER_ERROR,
173                        ErrorScope::Provider => StatusCode::SERVICE_UNAVAILABLE,
174                        ErrorScope::Egress | ErrorScope::Schema => StatusCode::BAD_GATEWAY,
175                    }
176                };
177                (status, json!({"error": e.to_string()}))
178            }
179        };
180        (status, Json(body)).into_response()
181    }
182}
183
184impl From<phrona::Error> for AppError {
185    fn from(e: phrona::Error) -> Self {
186        Self(ErrorKind::Internal(e))
187    }
188}
189
190type AppResult<T> = Result<T, AppError>;
191
192/// Query-string extractor whose rejection is a JSON 400 instead of axum's
193/// default plain-text response.
194pub struct JsonQuery<T>(pub T);
195
196impl<S, T: DeserializeOwned> FromRequestParts<S> for JsonQuery<T>
197where
198    S: Send + Sync,
199{
200    type Rejection = AppError;
201
202    async fn from_request_parts(
203        parts: &mut axum::http::request::Parts,
204        state: &S,
205    ) -> AppResult<Self> {
206        let q = Query::<T>::from_request_parts(parts, state)
207            .await
208            .map_err(|e| {
209                use std::error::Error as _;
210                let msg = match e {
211                    QueryRejection::FailedToDeserializeQueryString(e) => e
212                        .source()
213                        .map(|s| s.to_string())
214                        .unwrap_or_else(|| e.to_string()),
215                    other => other.to_string(),
216                };
217                AppError::bad_request(format!("invalid query parameters: {msg}"))
218            })?;
219        Ok(Self(q.0))
220    }
221}
222
223/// JSON body extractor whose rejection (missing body, malformed JSON,
224/// wrong content type) is a JSON 400 instead of axum's defaults. Bodies
225/// larger than the configured `server.max_body_bytes` are rejected with a
226/// 413 before deserialization.
227pub struct JsonBody<T>(pub T);
228
229impl<T: DeserializeOwned> FromRequest<Arc<AppState>> for JsonBody<T> {
230    type Rejection = AppError;
231
232    async fn from_request(req: Request, state: &Arc<AppState>) -> AppResult<Self> {
233        if let Some(len) = req
234            .headers()
235            .get(axum::http::header::CONTENT_LENGTH)
236            .and_then(|v| v.to_str().ok())
237            .and_then(|s| s.parse::<u64>().ok())
238            && len > state.max_body_bytes
239        {
240            return Err(AppError::body_too_large(state.max_body_bytes));
241        }
242        let Json(v) = Json::<T>::from_request(req, state)
243            .await
244            .map_err(|e| match &e {
245                JsonRejection::BytesRejection(BytesRejection::FailedToBufferBody(
246                    FailedToBufferBody::LengthLimitError(_),
247                )) => AppError::body_too_large(state.max_body_bytes),
248                _ => AppError::bad_request(format!("invalid JSON body: {e}")),
249            })?;
250        Ok(Self(v))
251    }
252}
253
254#[derive(Deserialize)]
255struct SearchParams {
256    q: String,
257    category: Option<String>,
258    engines: Option<String>,
259    page: Option<u32>,
260    max_results: Option<usize>,
261    safesearch: Option<String>,
262    region: Option<String>,
263    language: Option<String>,
264    time_range: Option<String>,
265    filters: Option<String>,
266}
267
268fn build_options(p: &SearchParams, max_results_limit: usize) -> AppResult<SearchOptions> {
269    let mut opts = SearchOptions::new(p.q.clone());
270    if let Some(c) = &p.category {
271        opts.category = c.parse::<Category>().map_err(|_| {
272            AppError::bad_request(format!(
273                "invalid category '{c}', expected one of: web, images, news, videos, books"
274            ))
275        })?;
276    }
277    if let Some(es) = &p.engines {
278        for name in es.split(',').map(str::trim).filter(|s| !s.is_empty()) {
279            if engine::engine_by_name(name).is_none() {
280                return Err(AppError::bad_request(format!(
281                    "unknown engine '{name}'. Available: {}",
282                    engine::list()
283                        .iter()
284                        .map(|e| e.name())
285                        .collect::<Vec<_>>()
286                        .join(", ")
287                )));
288            }
289            opts.engines.push(name.to_string());
290        }
291    }
292    if let Some(page) = p.page {
293        opts.page = page.max(1);
294    }
295    if let Some(m) = p.max_results {
296        opts.max_results = m.clamp(1, max_results_limit);
297    }
298    if let Some(s) = &p.safesearch {
299        opts.safesearch = s.parse::<phrona::SafeSearch>().map_err(|_| {
300            AppError::bad_request("invalid safesearch, expected off|moderate|strict")
301        })?;
302    }
303    if let Some(t) = &p.time_range {
304        opts.time_range = Some(t.parse::<TimeRange>().map_err(|_| {
305            AppError::bad_request("invalid time_range, expected day|week|month|year")
306        })?);
307    }
308    opts.region = p.region.clone();
309    opts.language = p.language.clone();
310    opts.filters = p.filters.clone();
311    Ok(opts)
312}
313
314fn header_key(headers: &HeaderMap) -> Option<String> {
315    headers
316        .get("x-api-key")
317        .and_then(|v| v.to_str().ok())
318        .map(str::to_string)
319        .or_else(|| {
320            headers
321                .get(axum::http::header::AUTHORIZATION)
322                .and_then(|v| v.to_str().ok())
323                .and_then(|v| v.strip_prefix("Bearer ").map(str::to_string))
324        })
325}
326
327/// Resolve the API key for endpoints that accept credentials in the JSON
328/// body (e.g. the Tavily-compatible routes): the body key wins, headers are
329/// the fallback. Query strings are never consulted.
330pub(crate) fn auth_key(headers: &HeaderMap, body_key: Option<&str>) -> Option<String> {
331    body_key.map(str::to_string).or_else(|| header_key(headers))
332}
333
334/// Header-only auth for GET endpoints. Query-string `api_key` parameters are
335/// rejected with a 400: credentials in URLs leak through logs, proxies and
336/// referrers and are a classic SSRF/credential-theft vector.
337pub struct HeaderAuth {
338    key: Option<String>,
339}
340
341impl HeaderAuth {
342    pub fn key(&self) -> Option<&str> {
343        self.key.as_deref()
344    }
345}
346
347impl<S> FromRequestParts<S> for HeaderAuth
348where
349    S: Send + Sync,
350{
351    type Rejection = AppError;
352
353    async fn from_request_parts(
354        parts: &mut axum::http::request::Parts,
355        _state: &S,
356    ) -> AppResult<Self> {
357        if let Some(query) = parts.uri.query()
358            && url::form_urlencoded::parse(query.as_bytes()).any(|(k, _)| k == "api_key")
359        {
360            return Err(AppError::bad_request(
361                "api_key in the query string is disallowed for security; use the x-api-key header or Authorization: Bearer instead",
362            ));
363        }
364        Ok(Self {
365            key: header_key(&parts.headers),
366        })
367    }
368}
369
370/// In-memory fixed-window rate limiter middleware: honors
371/// Resolve the client IP for rate limiting. The peer address is used
372/// directly unless it belongs to `trusted_proxies` (an operator-configured
373/// list of reverse proxies), in which case the leftmost `X-Forwarded-For`
374/// address is trusted — a proxy chain appends each hop, so the first entry
375/// is the original client. Without a trusted proxy, the header is never
376/// consulted, so it cannot be spoofed.
377fn client_ip(trusted: &[IpAddr], peer: SocketAddr, headers: &HeaderMap) -> IpAddr {
378    if trusted.contains(&peer.ip()) {
379        if let Some(ff) = headers
380            .get("x-forwarded-for")
381            .and_then(|v| v.to_str().ok())
382            .and_then(|v| v.split(',').next())
383        {
384            if let Ok(ip) = ff.trim().parse::<IpAddr>() {
385                return ip;
386            }
387        }
388    }
389    peer.ip()
390}
391
392/// `server.rate_limit_per_minute` per client IP (falling back to a single
393/// global bucket when the peer address is unknown, e.g. in tests).
394async fn rate_limit(
395    State(state): State<Arc<AppState>>,
396    req: Request,
397    next: Next,
398) -> AppResult<Response> {
399    let ip = req
400        .extensions()
401        .get::<ConnectInfo<SocketAddr>>()
402        .map(|c| client_ip(&state.trusted_proxies, c.0, req.headers()));
403    if !state.check_rate(ip) {
404        return Err(AppError::rate_limited(state.rate_limit_per_minute));
405    }
406    Ok(next.run(req).await)
407}
408
409async fn health(State(state): State<Arc<AppState>>) -> Json<Value> {
410    let web = engine::engines_for(Category::Web).len();
411    let images = engine::engines_for(Category::Images).len();
412    let news = engine::engines_for(Category::News).len();
413    let videos = engine::engines_for(Category::Videos).len();
414    let books = engine::engines_for(Category::Books).len();
415    Json(json!({
416        "status": "ok",
417        "version": phrona::version(),
418        "uptime_s": state.started.elapsed().as_secs(),
419        "engines": {"web": web, "images": images, "news": news, "videos": videos, "books": books},
420        "auth": state.api_key.is_some(),
421    }))
422}
423
424#[derive(Deserialize)]
425struct EnginesParams {
426    category: Option<String>,
427}
428
429async fn engines(JsonQuery(p): JsonQuery<EnginesParams>) -> AppResult<Json<Value>> {
430    let cats: Vec<Category> = match p.category.as_deref() {
431        Some(c) => vec![c.parse::<Category>().map_err(|_| {
432            AppError::bad_request(
433                "invalid category, expected one of: web, images, news, videos, books",
434            )
435        })?],
436        None => Category::ALL.to_vec(),
437    };
438    let mut out = serde_json::Map::new();
439    for cat in cats {
440        out.insert(
441            cat.as_str().to_string(),
442            json!(
443                phrona::available_engines(cat)
444                    .iter()
445                    .map(|e| e.name.clone())
446                    .collect::<Vec<_>>()
447            ),
448        );
449    }
450    Ok(Json(Value::Object(out)))
451}
452
453async fn search_route(
454    State(state): State<Arc<AppState>>,
455    auth: HeaderAuth,
456    JsonQuery(p): JsonQuery<SearchParams>,
457) -> AppResult<Json<SearchResponse>> {
458    if !state.authorized(auth.key()) {
459        return Err(AppError(ErrorKind::Unauthorized));
460    }
461    let opts = build_options(&p, state.max_results_limit)?;
462    let resp = state.client.search(opts).await?;
463    Ok(Json(resp))
464}
465
466#[derive(Deserialize)]
467struct SuggestParams {
468    q: String,
469    source: Option<String>,
470    region: Option<String>,
471}
472
473async fn suggest_route(
474    State(state): State<Arc<AppState>>,
475    auth: HeaderAuth,
476    JsonQuery(p): JsonQuery<SuggestParams>,
477) -> AppResult<Json<Value>> {
478    if !state.authorized(auth.key()) {
479        return Err(AppError(ErrorKind::Unauthorized));
480    }
481    let region = p.region.unwrap_or_else(|| "us-en".to_string());
482    match p.source.as_deref() {
483        Some(name) => {
484            let source = phrona::SuggestSource::from_name(name).ok_or_else(|| {
485                AppError::bad_request(format!(
486                    "unknown suggest source '{name}', expected one of: {}",
487                    phrona::SuggestSource::ALL
488                        .iter()
489                        .map(|s| s.name())
490                        .collect::<Vec<_>>()
491                        .join(", ")
492                ))
493            })?;
494            let suggestions = suggest(state.client.http(), source, &p.q, &region).await?;
495            Ok(Json(json!({
496                "query": p.q,
497                "source": name,
498                "suggestions": suggestions,
499            })))
500        }
501        None => {
502            let all = suggest_all(state.client.http(), &p.q, &region).await;
503            let suggestions: serde_json::Map<String, Value> = all
504                .into_iter()
505                .map(|(s, list)| (s.name().to_string(), json!(list)))
506                .collect();
507            Ok(Json(json!({
508                "query": p.q,
509                "suggestions": suggestions,
510            })))
511        }
512    }
513}
514
515/// Build the axum router from a [`PhronaConfig`]: the search client
516/// (profile / timeout / proxies / concurrency), the API key, the
517/// `max_results` clamp, the rate limit and the body-size cap all come from
518/// the config.
519///
520/// Protected endpoints (every `/v1/*` data route plus the Tavily-compatible
521/// aliases) run under the per-IP rate limiter; `GET /metrics`, `/health` and
522/// the frontend stay unauthenticated and unthrottled.
523pub fn router(cfg: PhronaConfig) -> Router {
524    let client = cfg
525        .search_client()
526        .expect("build search client")
527        .with_observer(Arc::new(metrics::EngineMetricsObserver));
528    let state = Arc::new(AppState::new(
529        client,
530        cfg.server.api_key.clone(),
531        cfg.max_results_limit(),
532        cfg.server.rate_limit_per_minute,
533        cfg.server.max_body_bytes,
534        cfg.server.trusted_proxies.clone(),
535    ));
536
537    let protected = Router::new()
538        .route("/v1/search", get(search_route))
539        .route("/v1/suggest", get(suggest_route))
540        .route(
541            "/v1/extract",
542            get(tools::extract_get).post(tools::extract_post),
543        )
544        .route("/v1/test", get(tools::test))
545        .route("/v1/grounding", get(grounding::get).post(grounding::post))
546        .route("/search", post(tavily::search))
547        .route("/v1/tavily", post(tavily::search))
548        .layer(middleware::from_fn_with_state(state.clone(), rate_limit));
549
550    protected
551        .merge(
552            Router::new()
553                .route("/", get(frontend::index))
554                .route("/health", get(health))
555                .route("/metrics", get(metrics::metrics_route))
556                .route("/v1/engines", get(engines))
557                .nest_service(
558                    "/static",
559                    tower_http::services::ServeDir::new(frontend::frontend_dir()),
560                )
561                .fallback(frontend::index),
562        )
563        .layer(DefaultBodyLimit::max(cfg.server.max_body_bytes as usize))
564        .layer(middleware::from_fn(metrics::http_layer))
565        .layer(CorsLayer::permissive())
566        .layer(TraceLayer::new_for_http())
567        .with_state(state)
568}
569
570/// Wait for Ctrl+C or SIGTERM, then return so the server can drain
571/// in-flight requests gracefully.
572pub async fn shutdown_signal() {
573    let ctrl_c = async {
574        tokio::signal::ctrl_c()
575            .await
576            .expect("install Ctrl+C handler");
577    };
578    #[cfg(unix)]
579    let terminate = async {
580        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
581            .expect("install SIGTERM handler")
582            .recv()
583            .await;
584    };
585    #[cfg(not(unix))]
586    let terminate = std::future::pending::<()>();
587    tokio::select! {
588        _ = ctrl_c => {}
589        _ = terminate => {}
590    }
591}
592
593/// Serve the REST API on `addr`. Blocks until the server stops (Ctrl+C or
594/// SIGTERM trigger a graceful shutdown that drains in-flight requests).
595pub async fn serve(addr: SocketAddr, cfg: PhronaConfig) -> anyhow::Result<()> {
596    let listener = tokio::net::TcpListener::bind(addr).await?;
597    tracing::info!("phrona-api listening on http://{addr}");
598    let app = router(cfg).into_make_service_with_connect_info::<SocketAddr>();
599    axum::serve(listener, app)
600        .with_graceful_shutdown(shutdown_signal())
601        .await?;
602    Ok(())
603}
604
605/// Default bind address when none is configured.
606pub fn default_addr() -> SocketAddr {
607    "127.0.0.1:8080".parse().expect("static addr")
608}
609
610#[cfg(test)]
611mod tests {
612    use super::*;
613    use axum::body::Body;
614    use axum::http::Request;
615
616    fn router_with(mut cfg: PhronaConfig) -> Router {
617        cfg.server.api_key = Some("test-secret".into());
618        super::router(cfg)
619    }
620
621    #[test]
622    fn app_error_maps_to_status_codes() {
623        for (status, kind) in [
624            (StatusCode::BAD_REQUEST, ErrorKind::BadRequest("x".into())),
625            (StatusCode::UNAUTHORIZED, ErrorKind::Unauthorized),
626            (
627                StatusCode::BAD_GATEWAY,
628                ErrorKind::Internal(phrona::Error::schema("e", "bad body")),
629            ),
630            (
631                StatusCode::TOO_MANY_REQUESTS,
632                ErrorKind::Internal(phrona::Error::rate_limited("e", None)),
633            ),
634            (
635                StatusCode::BAD_REQUEST,
636                ErrorKind::Internal(phrona::Error::invalid_query("e", "bad")),
637            ),
638            (
639                StatusCode::INTERNAL_SERVER_ERROR,
640                ErrorKind::Internal(phrona::Error::internal("e", "boom")),
641            ),
642        ] {
643            let resp = AppError(kind).into_response();
644            assert_eq!(resp.status(), status);
645        }
646    }
647
648    #[tokio::test]
649    async fn app_error_body_is_json() {
650        let resp = AppError(ErrorKind::BadRequest("bad input".into())).into_response();
651        let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
652        let v: Value = serde_json::from_slice(&body).unwrap();
653        assert_eq!(v["error"], "bad input");
654    }
655
656    #[test]
657    fn authorized_semantics() {
658        let client = PhronaConfig::defaults().search_client().unwrap();
659        let state = AppState::new(
660            client,
661            Some("test-secret".into()),
662            1000,
663            0,
664            100_000,
665            Vec::new(),
666        );
667        assert!(state.authorized(Some("test-secret")));
668        // wrong key, shorter key, longer key, missing key
669        assert!(!state.authorized(Some("wrong")));
670        assert!(!state.authorized(Some("test-secre")));
671        assert!(!state.authorized(Some("test-secret2")));
672        assert!(!state.authorized(None));
673        // no configured key: every request is accepted
674        let open = AppState::new(
675            PhronaConfig::defaults().search_client().unwrap(),
676            None,
677            1000,
678            0,
679            100_000,
680            Vec::new(),
681        );
682        assert!(open.authorized(None));
683        assert!(open.authorized(Some("anything")));
684    }
685
686    #[test]
687    fn client_ip_honors_trusted_proxies_only() {
688        let peer = "10.0.0.5:443".parse::<SocketAddr>().unwrap();
689        let mut headers = HeaderMap::new();
690        headers.insert(
691            "x-forwarded-for",
692            axum::http::HeaderValue::from_static("203.0.113.9, 10.0.0.1"),
693        );
694        // untrusted peer: the header is ignored (no spoofing)
695        assert_eq!(client_ip(&[], peer, &headers), peer.ip());
696        // trusted peer: leftmost forwarded address wins
697        assert_eq!(
698            client_ip(&["10.0.0.5".parse().unwrap()], peer, &headers),
699            "203.0.113.9".parse::<IpAddr>().unwrap()
700        );
701        // trusted peer with malformed header falls back to the peer
702        let mut bad = HeaderMap::new();
703        bad.insert(
704            "x-forwarded-for",
705            axum::http::HeaderValue::from_static("not-an-ip"),
706        );
707        assert_eq!(
708            client_ip(&["10.0.0.5".parse().unwrap()], peer, &bad),
709            peer.ip()
710        );
711        // a different trusted proxy listed but not the actual peer: no trust
712        assert_eq!(
713            client_ip(&["10.0.0.9".parse().unwrap()], peer, &headers),
714            peer.ip()
715        );
716    }
717
718    #[tokio::test]
719    async fn json_query_rejects_missing_fields_as_bad_request() {
720        use tower::ServiceExt;
721        let router = router_with(PhronaConfig::defaults());
722        let resp = router
723            .oneshot(
724                Request::builder()
725                    .uri("/v1/search")
726                    .header("x-api-key", "test-secret")
727                    .body(Body::empty())
728                    .unwrap(),
729            )
730            .await
731            .unwrap();
732        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
733        let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
734        let v: Value = serde_json::from_slice(&body).unwrap();
735        assert_eq!(
736            v["error"].as_str().unwrap().to_lowercase(),
737            "invalid query parameters: missing field `q`"
738        );
739    }
740
741    #[tokio::test]
742    async fn query_string_api_key_is_rejected_with_400() {
743        use tower::ServiceExt;
744        let router = router_with(PhronaConfig::defaults());
745        for path in [
746            "/v1/search?q=rust&api_key=test-secret",
747            "/v1/suggest?q=ru&api_key=test-secret",
748            "/v1/test?query=rust&api_key=test-secret",
749            "/v1/extract?url=https://example.com&api_key=test-secret",
750            "/v1/grounding?query=rust&api_key=test-secret",
751        ] {
752            let resp = router
753                .clone()
754                .oneshot(Request::builder().uri(path).body(Body::empty()).unwrap())
755                .await
756                .unwrap();
757            assert_eq!(resp.status(), StatusCode::BAD_REQUEST, "path {path}");
758            let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
759            let v: Value = serde_json::from_slice(&body).unwrap();
760            let msg = v["error"].as_str().unwrap_or("");
761            assert!(
762                msg.contains("query string"),
763                "path {path} got unexpected message: {msg}"
764            );
765        }
766    }
767
768    #[tokio::test]
769    async fn header_auth_grants_access_before_validation() {
770        use tower::ServiceExt;
771        let router = router_with(PhronaConfig::defaults());
772        let resp = router
773            .oneshot(
774                Request::builder()
775                    .uri("/v1/search?q=rust&category=bogus")
776                    .header("x-api-key", "test-secret")
777                    .body(Body::empty())
778                    .unwrap(),
779            )
780            .await
781            .unwrap();
782        // auth passed (no 401); validation then rejects the category
783        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
784        let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
785        let v: Value = serde_json::from_slice(&body).unwrap();
786        assert!(
787            v["error"]
788                .as_str()
789                .unwrap_or("")
790                .contains("invalid category")
791        );
792    }
793
794    #[tokio::test]
795    async fn bearer_token_is_accepted() {
796        use tower::ServiceExt;
797        let router = router_with(PhronaConfig::defaults());
798        let resp = router
799            .oneshot(
800                Request::builder()
801                    .uri("/v1/search?q=rust&category=bogus")
802                    .header("authorization", "Bearer test-secret")
803                    .body(Body::empty())
804                    .unwrap(),
805            )
806            .await
807            .unwrap();
808        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
809    }
810
811    #[test]
812    fn auth_key_resolves_body_then_headers() {
813        let mut headers = HeaderMap::new();
814        headers.insert("x-api-key", "header-key".parse().unwrap());
815        assert_eq!(auth_key(&headers, None).as_deref(), Some("header-key"));
816        assert_eq!(
817            auth_key(&headers, Some("body-key")).as_deref(),
818            Some("body-key")
819        );
820        assert_eq!(auth_key(&headers, Some("")), Some(String::new()));
821        let mut bearer = HeaderMap::new();
822        bearer.insert("authorization", "Bearer b-key".parse().unwrap());
823        assert_eq!(auth_key(&bearer, None).as_deref(), Some("b-key"));
824        assert_eq!(auth_key(&HeaderMap::new(), None), None);
825    }
826
827    #[tokio::test]
828    async fn rate_limit_returns_429_after_window_exhausted() {
829        use tower::ServiceExt;
830        let mut cfg = PhronaConfig::defaults();
831        cfg.server.rate_limit_per_minute = 2;
832        let router = router_with(cfg);
833        for _ in 0..2 {
834            let resp = router
835                .clone()
836                .oneshot(
837                    Request::builder()
838                        .uri("/v1/search?q=rust&category=bogus")
839                        .body(Body::empty())
840                        .unwrap(),
841                )
842                .await
843                .unwrap();
844            // auth is checked before the search runs, so no network happens
845            assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
846        }
847        let resp = router
848            .clone()
849            .oneshot(
850                Request::builder()
851                    .uri("/v1/search?q=rust")
852                    .body(Body::empty())
853                    .unwrap(),
854            )
855            .await
856            .unwrap();
857        assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
858        let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
859        let v: Value = serde_json::from_slice(&body).unwrap();
860        assert!(v["error"].as_str().unwrap().contains("rate limit"));
861    }
862
863    #[tokio::test]
864    async fn oversized_body_is_413_json() {
865        use tower::ServiceExt;
866        let mut cfg = PhronaConfig::defaults();
867        cfg.server.max_body_bytes = 8;
868        let router = super::router(cfg);
869        let resp = router
870            .oneshot(
871                Request::builder()
872                    .method("POST")
873                    .uri("/v1/extract")
874                    .header("content-type", "application/json")
875                    .header("content-length", "30")
876                    .body(Body::from(r#"{"url": "https://example.com"}"#))
877                    .unwrap(),
878            )
879            .await
880            .unwrap();
881        assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE);
882        let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
883        let v: Value = serde_json::from_slice(&body).unwrap();
884        assert!(v["error"].as_str().unwrap().contains("8-byte limit"));
885    }
886
887    #[tokio::test]
888    async fn rate_limit_honors_connect_info_ip() {
889        use tower::ServiceExt;
890        let mut cfg = PhronaConfig::defaults();
891        cfg.server.api_key = Some("test-secret".into());
892        cfg.server.rate_limit_per_minute = 1;
893        let router = super::router(cfg);
894        // Same IP, different ephemeral ports (as real TCP clients produce):
895        // the window must be keyed on the IP, not the socket address, or a
896        // fresh bucket would be allocated per connection and the limiter
897        // would never trip.
898        let req = |port: u16| {
899            Request::builder()
900                .uri("/v1/search?q=rust&category=bogus")
901                .extension(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], port))))
902                .header("x-api-key", "test-secret")
903                .body(Body::empty())
904                .unwrap()
905        };
906        assert_eq!(
907            router.clone().oneshot(req(4242)).await.unwrap().status(),
908            StatusCode::BAD_REQUEST
909        );
910        assert_eq!(
911            router.clone().oneshot(req(4243)).await.unwrap().status(),
912            StatusCode::TOO_MANY_REQUESTS
913        );
914    }
915}