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
33struct 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 pub max_results_limit: usize,
45 pub rate_limit_per_minute: u32,
47 pub max_body_bytes: u64,
49 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 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 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 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
192pub 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
223pub 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
327pub(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
334pub 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
370fn 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
392async 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, ®ion).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, ®ion).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
515pub 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
570pub 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
593pub 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
605pub 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 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 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 assert_eq!(client_ip(&[], peer, &headers), peer.ip());
696 assert_eq!(
698 client_ip(&["10.0.0.5".parse().unwrap()], peer, &headers),
699 "203.0.113.9".parse::<IpAddr>().unwrap()
700 );
701 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 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 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 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 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}