Skip to main content

phrona_api/
lib.rs

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