1mod health;
4mod helpers;
5mod kbs;
6mod llms;
7mod objects;
8
9use crate::middleware::auth_middleware;
10use crate::state::AppState;
11use axum::Router;
12use axum::extract::{DefaultBodyLimit, Request};
13use axum::handler::Handler;
14use axum::http::HeaderName;
15use axum::middleware::from_fn_with_state;
16use axum::routing::get;
17use tower::ServiceBuilder;
18use tower_http::request_id::{
19 MakeRequestId, PropagateRequestIdLayer, RequestId, SetRequestIdLayer,
20};
21use tower_http::trace::TraceLayer;
22use uuid::Uuid;
23
24pub(crate) use helpers::lookup_kb;
26
27use health::{healthz, readyz};
28use kbs::{list_kbs, list_objects};
29use llms::llms_txt;
30use objects::{delete_object, get_object, head_object, patch_object, post_object, put_object};
31
32pub const MAX_BODY_BYTES: u64 = 16 * 1024 * 1024;
34
35#[derive(Clone, Copy, Default)]
37pub struct MakeRequestUuidV7;
38
39impl MakeRequestId for MakeRequestUuidV7 {
40 fn make_request_id<B>(&mut self, _req: &Request<B>) -> Option<RequestId> {
41 let id = Uuid::now_v7().to_string();
42 let hv = id.parse().ok()?;
43 Some(RequestId::new(hv))
44 }
45}
46
47pub fn build_router(state: AppState) -> Router {
49 let request_id_header = HeaderName::from_static("x-request-id");
50 Router::new()
51 .route("/healthz", get(healthz))
52 .route("/readyz", get(readyz))
53 .route("/llms.txt", get(llms_txt))
54 .route("/v1/knowledgebases", get(list_kbs))
55 .route("/v1/knowledgebases/{kb_slug}", get(list_objects))
56 .route(
57 "/v1/knowledgebases/{kb_slug}/search",
58 axum::routing::post(crate::search_route::search_kb).layer(
59 axum::extract::DefaultBodyLimit::max(crate::search_route::SEARCH_BODY_MAX_BYTES),
60 ),
61 )
62 .route(
63 "/v1/knowledgebases/{kb_slug}/{*object_path}",
64 get(get_object)
65 .head(head_object)
66 .put(put_object)
67 .delete(delete_object)
68 .patch(patch_object.layer(DefaultBodyLimit::disable()))
69 .post(post_object.layer(DefaultBodyLimit::disable())),
70 )
71 .layer(
72 ServiceBuilder::new()
73 .layer(DefaultBodyLimit::max(helpers::body_limit_usize(
74 MAX_BODY_BYTES,
75 )))
76 .layer(SetRequestIdLayer::new(
77 request_id_header.clone(),
78 MakeRequestUuidV7,
79 ))
80 .layer(PropagateRequestIdLayer::new(request_id_header))
81 .layer(TraceLayer::new_for_http())
82 .layer(from_fn_with_state(state.clone(), auth_middleware)),
83 )
84 .with_state(state)
85}
86
87#[cfg(test)]
88mod patch_route {
89 use super::*;
90 use async_trait::async_trait;
91 use axum::body::{Body, to_bytes};
92 use axum::http::StatusCode;
93 use axum::response::Response;
94 use bytes::Bytes;
95 use notedthat_core::{
96 ByteRange, ConditionalHeaders, KbManifest, KbSlug, ListResponse, ObjectMeta, ObjectPath,
97 ObjectRead, PutOutcome, Storage, StorageError,
98 };
99 use notedthat_indexer::IndexEvent;
100 use std::collections::BTreeMap;
101 use std::sync::Arc;
102 use tower::util::ServiceExt;
103
104 const KB: &str = "notes";
105 const OBJECT_PATH: &str = "patch.md";
106 const TOKEN: &str = "test-token-abc";
107
108 async fn router_with_object(
109 initial_body: &'static [u8],
110 max_patchable_size: u64,
111 ) -> (axum::Router, String) {
112 let kb = KbSlug::try_new(KB).unwrap();
113 let object_path = ObjectPath::try_from_str(OBJECT_PATH).unwrap();
114 let storage = Arc::new(crate::testing::InMemoryStorage::default());
115 let outcome = storage
116 .put_object(
117 &kb,
118 &object_path,
119 Bytes::from_static(initial_body),
120 Some("text/markdown"),
121 ConditionalHeaders::default(),
122 )
123 .await
124 .unwrap();
125
126 let (indexer_tx, _rx) = tokio::sync::mpsc::channel(16);
127 let mut kbs = BTreeMap::new();
128 kbs.insert(KB.to_string(), kb);
129 let router = build_router(AppState {
130 storage,
131 declared_kbs: Arc::new(kbs),
132 public_read_policies: Arc::new(BTreeMap::new()),
133 bearer_token: Arc::new(TOKEN.to_string()),
134 max_body_size: MAX_BODY_BYTES,
135 max_patchable_size,
136 indexer_tx,
137 searcher: Arc::new(crate::testing::NoopSearcher),
138 });
139
140 (router, outcome.etag.unwrap())
141 }
142
143 async fn object_with_etag(
144 storage: &crate::testing::InMemoryStorage,
145 kb: &KbSlug,
146 body: &'static [u8],
147 ) -> String {
148 storage
149 .put_object(
150 kb,
151 &ObjectPath::try_from_str(OBJECT_PATH).unwrap(),
152 Bytes::from_static(body),
153 Some("text/markdown"),
154 ConditionalHeaders::default(),
155 )
156 .await
157 .unwrap()
158 .etag
159 .unwrap()
160 }
161
162 fn router_with_storage(
163 storage: Arc<dyn Storage>,
164 kb: KbSlug,
165 max_patchable_size: u64,
166 indexer_tx: tokio::sync::mpsc::Sender<IndexEvent>,
167 ) -> axum::Router {
168 let mut kbs = BTreeMap::new();
169 kbs.insert(KB.to_string(), kb);
170 build_router(AppState {
171 storage,
172 declared_kbs: Arc::new(kbs),
173 public_read_policies: Arc::new(BTreeMap::new()),
174 bearer_token: Arc::new(TOKEN.to_string()),
175 max_body_size: MAX_BODY_BYTES,
176 max_patchable_size,
177 indexer_tx,
178 searcher: Arc::new(crate::testing::NoopSearcher),
179 })
180 }
181
182 async fn patch_request(
183 router: axum::Router,
184 header_name: &'static str,
185 header_value: &str,
186 if_match: Option<&str>,
187 body: Bytes,
188 ) -> Response {
189 let mut builder = Request::builder()
190 .method("PATCH")
191 .uri(format!("/v1/knowledgebases/{KB}/{OBJECT_PATH}"))
192 .header("authorization", format!("Bearer {TOKEN}"))
193 .header(header_name, header_value);
194 if let Some(etag) = if_match {
195 builder = builder.header(axum::http::header::IF_MATCH, etag);
196 }
197
198 router
199 .oneshot(builder.body(Body::from(body)).unwrap())
200 .await
201 .unwrap()
202 }
203
204 async fn assert_error_code(response: Response, expected_status: StatusCode, expected: &str) {
205 assert_eq!(response.status(), expected_status);
206 let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
207 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
208 assert_eq!(json["error"], expected);
209 }
210
211 #[tokio::test]
212 async fn bytes_content_range_returns_ok_with_etag_and_location() {
213 let (router, etag) = router_with_object(b"0123456789abcdefghij", MAX_BODY_BYTES).await;
214
215 let response = patch_request(
216 router,
217 "content-range",
218 "bytes 0-9/*",
219 Some(&etag),
220 Bytes::from_static(b"ABCDEFGHIJ"),
221 )
222 .await;
223
224 assert_eq!(response.status(), StatusCode::OK);
225 assert!(response.headers().get(axum::http::header::ETAG).is_some());
226 assert_eq!(
227 response
228 .headers()
229 .get(axum::http::header::LOCATION)
230 .unwrap(),
231 &format!("/v1/knowledgebases/{KB}/{OBJECT_PATH}")
232 );
233 assert!(
234 response
235 .headers()
236 .get(axum::http::header::CONTENT_RANGE)
237 .is_none()
238 );
239 assert!(response.headers().get("nt-patch-mode").is_none());
240 }
241
242 #[tokio::test]
243 async fn lines_content_range_returns_ok() {
244 let (router, etag) = router_with_object(b"one\ntwo\nthree\nfour\n", MAX_BODY_BYTES).await;
245
246 let response = patch_request(
247 router,
248 "content-range",
249 "lines 2-3/*",
250 Some(&etag),
251 Bytes::from_static(b"TWO\nTHREE\n"),
252 )
253 .await;
254
255 assert_eq!(response.status(), StatusCode::OK);
256 }
257
258 #[tokio::test]
259 async fn append_mode_without_if_match_returns_ok() {
260 let (router, _etag) = router_with_object(b"one\n", MAX_BODY_BYTES).await;
261
262 let response = patch_request(
263 router,
264 "nt-patch-mode",
265 "append",
266 None,
267 Bytes::from_static(b"two\n"),
268 )
269 .await;
270
271 assert_eq!(response.status(), StatusCode::OK);
272 }
273
274 #[tokio::test]
275 async fn append_mode_with_if_match_returns_ok() {
276 let (router, etag) = router_with_object(b"one\n", MAX_BODY_BYTES).await;
277
278 let response = patch_request(
279 router,
280 "nt-patch-mode",
281 "append",
282 Some(&etag),
283 Bytes::from_static(b"two\n"),
284 )
285 .await;
286
287 assert_eq!(response.status(), StatusCode::OK);
288 }
289
290 #[tokio::test]
291 async fn bytes_content_range_without_if_match_returns_invalid_request() {
292 let (router, _etag) = router_with_object(b"0123456789", MAX_BODY_BYTES).await;
293
294 let response = patch_request(
295 router,
296 "content-range",
297 "bytes 0-1/*",
298 None,
299 Bytes::from_static(b"AB"),
300 )
301 .await;
302
303 assert_error_code(response, StatusCode::BAD_REQUEST, "invalid_request").await;
304 }
305
306 #[tokio::test]
307 async fn if_match_star_returns_invalid_request() {
308 let (router, _etag) = router_with_object(b"0123456789", MAX_BODY_BYTES).await;
309
310 let response = patch_request(
311 router,
312 "content-range",
313 "bytes 0-1/*",
314 Some("*"),
315 Bytes::from_static(b"AB"),
316 )
317 .await;
318
319 assert_error_code(response, StatusCode::BAD_REQUEST, "invalid_request").await;
320 }
321
322 #[tokio::test]
323 async fn multi_value_if_match_returns_invalid_request() {
324 let (router, etag) = router_with_object(b"0123456789", MAX_BODY_BYTES).await;
325
326 let response = patch_request(
327 router,
328 "content-range",
329 "bytes 0-1/*",
330 Some(&format!("{etag}, \"other\"")),
331 Bytes::from_static(b"AB"),
332 )
333 .await;
334
335 assert_error_code(response, StatusCode::BAD_REQUEST, "invalid_request").await;
336 }
337
338 #[tokio::test]
339 async fn nonexistent_object_returns_not_found() {
340 let (router, etag) = router_with_object(b"0123456789", MAX_BODY_BYTES).await;
341
342 let response = router
343 .oneshot(
344 Request::builder()
345 .method("PATCH")
346 .uri(format!("/v1/knowledgebases/{KB}/missing.md"))
347 .header("authorization", format!("Bearer {TOKEN}"))
348 .header("content-range", "bytes 0-1/*")
349 .header(axum::http::header::IF_MATCH, etag)
350 .body(Body::from(Bytes::from_static(b"AB")))
351 .unwrap(),
352 )
353 .await
354 .unwrap();
355
356 assert_error_code(response, StatusCode::NOT_FOUND, "not_found").await;
357 }
358
359 #[tokio::test]
360 async fn body_larger_than_max_patchable_size_returns_payload_too_large() {
361 let (router, _etag) = router_with_object(b"one\n", 4).await;
362
363 let response = patch_request(
364 router,
365 "nt-patch-mode",
366 "append",
367 None,
368 Bytes::from_static(b"abcde"),
369 )
370 .await;
371
372 assert_error_code(response, StatusCode::PAYLOAD_TOO_LARGE, "payload_too_large").await;
373 }
374
375 mod errors {
376 use super::*;
377
378 #[derive(Clone)]
379 struct PutPreconditionFailedStorage {
380 inner: crate::testing::InMemoryStorage,
381 }
382
383 #[async_trait]
384 impl Storage for PutPreconditionFailedStorage {
385 async fn ensure_bucket(&self, kb: &KbSlug) -> Result<(), StorageError> {
386 self.inner.ensure_bucket(kb).await
387 }
388
389 async fn read_manifest(&self, kb: &KbSlug) -> Result<KbManifest, StorageError> {
390 self.inner.read_manifest(kb).await
391 }
392
393 async fn write_manifest(
394 &self,
395 kb: &KbSlug,
396 manifest: &KbManifest,
397 ) -> Result<(), StorageError> {
398 self.inner.write_manifest(kb, manifest).await
399 }
400
401 async fn head_object(
402 &self,
403 kb: &KbSlug,
404 path: &ObjectPath,
405 conditionals: ConditionalHeaders,
406 ) -> Result<ObjectMeta, StorageError> {
407 self.inner.head_object(kb, path, conditionals).await
408 }
409
410 async fn get_object(
411 &self,
412 kb: &KbSlug,
413 path: &ObjectPath,
414 range: Option<Vec<ByteRange>>,
415 conditionals: ConditionalHeaders,
416 ) -> Result<ObjectRead, StorageError> {
417 self.inner.get_object(kb, path, range, conditionals).await
418 }
419
420 async fn get_object_stream(
421 &self,
422 kb: &KbSlug,
423 path: &ObjectPath,
424 range: Option<Vec<ByteRange>>,
425 conditionals: ConditionalHeaders,
426 ) -> Result<notedthat_core::ObjectStream, StorageError> {
427 self.inner
428 .get_object_stream(kb, path, range, conditionals)
429 .await
430 }
431
432 async fn put_object(
433 &self,
434 _kb: &KbSlug,
435 _path: &ObjectPath,
436 _bytes: Bytes,
437 _content_type: Option<&str>,
438 _conditionals: ConditionalHeaders,
439 ) -> Result<PutOutcome, StorageError> {
440 Err(StorageError::PreconditionFailed)
441 }
442
443 async fn put_staged_object(
444 &self,
445 _kb: &KbSlug,
446 _path: &ObjectPath,
447 _body: notedthat_core::StagedBody,
448 _content_type: Option<&str>,
449 _conditionals: ConditionalHeaders,
450 ) -> Result<PutOutcome, StorageError> {
451 Err(StorageError::PreconditionFailed)
452 }
453
454 async fn copy_object(
455 &self,
456 kb: &KbSlug,
457 source: &ObjectPath,
458 destination: &ObjectPath,
459 options: notedthat_core::CopyObjectOptions,
460 ) -> Result<PutOutcome, StorageError> {
461 self.inner
462 .copy_object(kb, source, destination, options)
463 .await
464 }
465
466 async fn delete_object(
467 &self,
468 kb: &KbSlug,
469 path: &ObjectPath,
470 conditionals: ConditionalHeaders,
471 ) -> Result<(), StorageError> {
472 self.inner.delete_object(kb, path, conditionals).await
473 }
474
475 async fn list_objects(
476 &self,
477 kb: &KbSlug,
478 prefix: Option<&str>,
479 limit: u32,
480 cursor: Option<&str>,
481 ) -> Result<ListResponse, StorageError> {
482 self.inner.list_objects(kb, prefix, limit, cursor).await
483 }
484 }
485
486 #[tokio::test]
487 async fn missing_if_match_for_bytes_mode_returns_invalid_request() {
488 let (router, _etag) = router_with_object(b"0123456789", MAX_BODY_BYTES).await;
489
490 let response = patch_request(
491 router,
492 "content-range",
493 "bytes 0-9/*",
494 None,
495 Bytes::from_static(b"ABCDEFGHIJ"),
496 )
497 .await;
498
499 assert_error_code(response, StatusCode::BAD_REQUEST, "invalid_request").await;
500 }
501
502 #[tokio::test]
503 async fn body_larger_than_max_patchable_size_returns_payload_too_large() {
504 let (router, _etag) = router_with_object(b"one\n", 10).await;
505
506 let response = patch_request(
507 router,
508 "nt-patch-mode",
509 "append",
510 None,
511 Bytes::from_static(b"more than ten bytes"),
512 )
513 .await;
514
515 assert_error_code(response, StatusCode::PAYLOAD_TOO_LARGE, "payload_too_large").await;
516 }
517
518 #[tokio::test]
519 async fn pre_splice_object_larger_than_max_patchable_size_returns_payload_too_large() {
520 let (router, _etag) = router_with_object(b"already too large", 10).await;
521
522 let response = patch_request(
523 router,
524 "nt-patch-mode",
525 "append",
526 None,
527 Bytes::from_static(b"!"),
528 )
529 .await;
530
531 assert_error_code(response, StatusCode::PAYLOAD_TOO_LARGE, "payload_too_large").await;
532 }
533
534 #[tokio::test]
535 async fn post_splice_body_larger_than_max_patchable_size_returns_payload_too_large() {
536 let (router, _etag) = router_with_object(b"123456", 10).await;
537
538 let response = patch_request(
539 router,
540 "nt-patch-mode",
541 "append",
542 None,
543 Bytes::from_static(b"78901"),
544 )
545 .await;
546
547 assert_error_code(response, StatusCode::PAYLOAD_TOO_LARGE, "payload_too_large").await;
548 }
549
550 #[tokio::test]
551 async fn line_range_out_of_bounds_returns_dual_416_headers_and_empty_body() {
552 let (router, etag) =
553 router_with_object(b"one\ntwo\nthree\nfour\nfive\n", MAX_BODY_BYTES).await;
554
555 let response = patch_request(
556 router,
557 "content-range",
558 "lines 100-200/*",
559 Some(&etag),
560 Bytes::from_static(b"replacement\n"),
561 )
562 .await;
563
564 assert_eq!(response.status(), StatusCode::RANGE_NOT_SATISFIABLE);
565 assert_eq!(
566 response.headers().get("content-range").unwrap(),
567 "lines */5"
568 );
569 assert_eq!(
570 response.headers().get("x-content-range-bytes").unwrap(),
571 "*/24"
572 );
573 let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
574 assert!(body.is_empty());
575 }
576
577 #[tokio::test]
578 async fn if_match_mismatch_after_retries_returns_precondition_failed_without_content_range()
579 {
580 let kb = KbSlug::try_new(KB).unwrap();
581 let inner = crate::testing::InMemoryStorage::default();
582 let etag = object_with_etag(&inner, &kb, b"0123456789").await;
583 let (indexer_tx, _rx) = tokio::sync::mpsc::channel(16);
584 let router = router_with_storage(
585 Arc::new(PutPreconditionFailedStorage { inner }),
586 kb,
587 MAX_BODY_BYTES,
588 indexer_tx,
589 );
590
591 let response = patch_request(
592 router,
593 "content-range",
594 "bytes 0-1/*",
595 Some(&etag),
596 Bytes::from_static(b"AB"),
597 )
598 .await;
599
600 assert_eq!(response.status(), StatusCode::PRECONDITION_FAILED);
601 assert!(
602 response
603 .headers()
604 .get(axum::http::header::CONTENT_RANGE)
605 .is_none()
606 );
607 let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
608 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
609 assert_eq!(json["error"], "precondition_failed");
610 }
611
612 #[tokio::test]
613 async fn indexer_queue_full_returns_backend_unavailable_with_retry_after() {
614 let kb = KbSlug::try_new(KB).unwrap();
615 let storage = crate::testing::InMemoryStorage::default();
616 let etag = object_with_etag(&storage, &kb, b"one\n").await;
617 let (indexer_tx, _rx) = tokio::sync::mpsc::channel(1);
618 indexer_tx
619 .try_send(IndexEvent::Upsert {
620 kb: kb.clone(),
621 object_key: ObjectPath::try_from_str("queued.md").unwrap(),
622 etag: "queued".to_string(),
623 mtime: 0,
624 })
625 .unwrap();
626 let router = router_with_storage(Arc::new(storage), kb, MAX_BODY_BYTES, indexer_tx);
627
628 let response = patch_request(
629 router,
630 "nt-patch-mode",
631 "append",
632 Some(&etag),
633 Bytes::from_static(b"two\n"),
634 )
635 .await;
636
637 assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
638 assert_eq!(response.headers().get("retry-after").unwrap(), "5");
639 let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
640 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
641 assert_eq!(json["error"], "backend_unavailable");
642 }
643 }
644}
645
646#[cfg(test)]
647mod line_range_get {
648 use super::*;
649 use axum::body::{Body, to_bytes};
650 use axum::http::StatusCode;
651 use axum::response::Response;
652 use bytes::Bytes;
653 use notedthat_core::{ConditionalHeaders, KbSlug, ObjectPath, Storage};
654 use std::collections::BTreeMap;
655 use std::sync::Arc;
656 use tower::util::ServiceExt;
657
658 const KB: &str = "notes";
659 const TOKEN: &str = "test-token-abc";
660
661 fn twenty_line_markdown() -> String {
662 markdown_lines(1, 20)
663 }
664
665 fn markdown_lines(first: u32, last: u32) -> String {
666 let mut body = String::new();
667 for line in first..=last {
668 std::fmt::Write::write_fmt(&mut body, format_args!("line {line:02}\n")).unwrap();
669 }
670 body
671 }
672
673 async fn router_with_markdown_object(body: String) -> axum::Router {
674 let kb = KbSlug::try_new(KB).unwrap();
675 let object_path = ObjectPath::try_from_str("ranges.md").unwrap();
676 let storage = Arc::new(crate::testing::InMemoryStorage::default());
677 storage
678 .put_object(
679 &kb,
680 &object_path,
681 Bytes::from(body),
682 Some("text/markdown"),
683 ConditionalHeaders::default(),
684 )
685 .await
686 .unwrap();
687
688 let (indexer_tx, _rx) = tokio::sync::mpsc::channel(1);
689 let mut kbs = BTreeMap::new();
690 kbs.insert(KB.to_string(), kb);
691 build_router(AppState {
692 storage,
693 declared_kbs: Arc::new(kbs),
694 public_read_policies: Arc::new(BTreeMap::new()),
695 bearer_token: Arc::new(TOKEN.to_string()),
696 max_body_size: MAX_BODY_BYTES,
697 max_patchable_size: MAX_BODY_BYTES,
698 indexer_tx,
699 searcher: Arc::new(crate::testing::NoopSearcher),
700 })
701 }
702
703 async fn get_ranges_md(router: axum::Router, range: &str) -> Response {
704 router
705 .oneshot(
706 Request::builder()
707 .method("GET")
708 .uri(format!("/v1/knowledgebases/{KB}/ranges.md"))
709 .header("authorization", format!("Bearer {TOKEN}"))
710 .header(axum::http::header::RANGE, range)
711 .body(Body::empty())
712 .unwrap(),
713 )
714 .await
715 .unwrap()
716 }
717
718 #[tokio::test]
719 async fn returns_first_five_lines_when_closed_range_requested() {
720 let router = router_with_markdown_object(twenty_line_markdown()).await;
721
722 let response = get_ranges_md(router, "lines=1-5").await;
723
724 assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT);
725 let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
726 assert_eq!(body, Bytes::from(markdown_lines(1, 5)));
727 }
728
729 #[tokio::test]
730 async fn returns_last_three_lines_when_suffix_range_requested() {
731 let router = router_with_markdown_object(twenty_line_markdown()).await;
732
733 let response = get_ranges_md(router, "lines=-3").await;
734
735 assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT);
736 let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
737 assert_eq!(body, Bytes::from(markdown_lines(18, 20)));
738 }
739
740 #[tokio::test]
741 async fn returns_empty_body_when_insert_range_requested() {
742 let router = router_with_markdown_object(twenty_line_markdown()).await;
743
744 let response = get_ranges_md(router, "lines=5-4").await;
745
746 assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT);
747 let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
748 assert!(body.is_empty());
749 }
750
751 #[tokio::test]
752 async fn returns_full_body_when_unknown_range_unit_requested() {
753 let body = twenty_line_markdown();
754 let router = router_with_markdown_object(body.clone()).await;
755
756 let response = get_ranges_md(router, "items=0-5").await;
757
758 assert_eq!(response.status(), StatusCode::OK);
759 let actual = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
760 assert_eq!(actual, Bytes::from(body));
761 }
762
763 mod headers {
764 use super::*;
765
766 fn ten_line_markdown() -> String {
767 markdown_lines(1, 10)
768 }
769
770 #[tokio::test]
771 async fn returns_line_and_byte_content_ranges_when_closed_range_requested() {
772 let router = router_with_markdown_object(ten_line_markdown()).await;
773
774 let response = get_ranges_md(router, "lines=2-4").await;
775
776 assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT);
777 assert_eq!(
778 response.headers().get("Content-Range").unwrap(),
779 "lines 2-4/10"
780 );
781 assert_eq!(
782 response.headers().get("X-Content-Range-Bytes").unwrap(),
783 "8-31/80"
784 );
785 }
786
787 #[tokio::test]
788 async fn returns_slice_content_length_when_closed_range_requested() {
789 let router = router_with_markdown_object(ten_line_markdown()).await;
790
791 let response = get_ranges_md(router, "lines=2-4").await;
792
793 assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT);
794 assert_eq!(response.headers().get("content-length").unwrap(), "24");
795 }
796
797 #[tokio::test]
798 async fn returns_zero_length_and_empty_byte_range_when_insert_range_requested() {
799 let router = router_with_markdown_object(ten_line_markdown()).await;
800
801 let response = get_ranges_md(router, "lines=5-4").await;
802
803 assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT);
804 assert_eq!(response.headers().get("content-length").unwrap(), "0");
805 assert_eq!(
806 response.headers().get("Content-Range").unwrap(),
807 "lines 5-4/10"
808 );
809 assert_eq!(
810 response.headers().get("X-Content-Range-Bytes").unwrap(),
811 "32-31/80"
812 );
813 }
814
815 #[tokio::test]
816 async fn omits_line_byte_range_header_when_byte_range_requested() {
817 let router = router_with_markdown_object(ten_line_markdown()).await;
818
819 let response = get_ranges_md(router, "bytes=0-9").await;
820
821 assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT);
822 assert!(response.headers().get("X-Content-Range-Bytes").is_none());
823 }
824 }
825}
826
827#[cfg(test)]
828mod tests {
829 use super::*;
830 use axum::body::{Body, to_bytes};
831 use axum::http::StatusCode;
832 use axum::response::Response;
833 use bytes::Bytes;
834 use notedthat_core::{ConditionalHeaders, KbSlug, ObjectPath, Storage, StorageError};
835 use notedthat_indexer::IndexEvent;
836 use std::collections::BTreeMap;
837 use std::sync::Arc;
838 use tower::util::ServiceExt;
839
840 const KB: &str = "notes";
841 const TOKEN: &str = "test-token-abc";
842
843 fn router() -> axum::Router {
844 let kb = KbSlug::try_new(KB).unwrap();
845 let mut kbs = BTreeMap::new();
846 kbs.insert(KB.to_string(), kb);
847 let (indexer_tx, mut rx) = tokio::sync::mpsc::channel(16);
848 tokio::spawn(async move { while rx.recv().await.is_some() {} });
849
850 build_router(AppState {
851 storage: Arc::new(crate::testing::InMemoryStorage::default()),
852 declared_kbs: Arc::new(kbs),
853 public_read_policies: Arc::new(BTreeMap::new()),
854 bearer_token: Arc::new(TOKEN.to_string()),
855 max_body_size: MAX_BODY_BYTES,
856 max_patchable_size: MAX_BODY_BYTES,
857 indexer_tx,
858 searcher: Arc::new(crate::testing::NoopSearcher),
859 })
860 }
861
862 async fn put_object(router: axum::Router, path: &str, body: &'static [u8]) -> Response {
863 router
864 .oneshot(
865 Request::builder()
866 .method("PUT")
867 .uri(format!("/v1/knowledgebases/{KB}/{path}"))
868 .header("authorization", format!("Bearer {TOKEN}"))
869 .header(axum::http::header::CONTENT_TYPE, "text/markdown")
870 .body(Body::from(Bytes::from_static(body)))
871 .unwrap(),
872 )
873 .await
874 .unwrap()
875 }
876
877 async fn put_object_etag(router: axum::Router, path: &str, body: &'static [u8]) -> String {
878 let response = put_object(router, path, body).await;
879 assert_eq!(response.status(), StatusCode::CREATED);
880 response
881 .headers()
882 .get(axum::http::header::ETAG)
883 .unwrap()
884 .to_str()
885 .unwrap()
886 .to_string()
887 }
888
889 async fn get_object(router: axum::Router, path: &str) -> Response {
890 router
891 .oneshot(
892 Request::builder()
893 .method("GET")
894 .uri(format!("/v1/knowledgebases/{KB}/{path}"))
895 .header("authorization", format!("Bearer {TOKEN}"))
896 .body(Body::empty())
897 .unwrap(),
898 )
899 .await
900 .unwrap()
901 }
902
903 async fn post_replace(
904 router: axum::Router,
905 path: &str,
906 if_match: &str,
907 body: &'static [u8],
908 ) -> Response {
909 router
910 .oneshot(
911 Request::builder()
912 .method("POST")
913 .uri(format!("/v1/knowledgebases/{KB}/replace/{path}"))
914 .header("authorization", format!("Bearer {TOKEN}"))
915 .header(axum::http::header::CONTENT_TYPE, "application/json")
916 .header(axum::http::header::IF_MATCH, if_match)
917 .body(Body::from(Bytes::from_static(body)))
918 .unwrap(),
919 )
920 .await
921 .unwrap()
922 }
923
924 #[tokio::test]
925 async fn get_on_replace_prefixed_path_still_reads_object_via_catch_all() {
926 let router = router();
927 put_object_etag(router.clone(), "replace/foo.md", b"hi").await;
928
929 let response = get_object(router, "replace/foo.md").await;
930
931 assert_eq!(response.status(), StatusCode::OK);
932 let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
933 assert_eq!(&body[..], b"hi");
934 }
935
936 #[tokio::test]
937 async fn patch_on_replace_prefixed_path_still_reaches_patch_object() {
938 let router = router();
939 let etag = put_object_etag(router.clone(), "replace/bar.md", b"old\n").await;
940
941 let response = router
942 .oneshot(
943 Request::builder()
944 .method("PATCH")
945 .uri(format!("/v1/knowledgebases/{KB}/replace/bar.md"))
946 .header("authorization", format!("Bearer {TOKEN}"))
947 .header(axum::http::header::CONTENT_RANGE, "lines 1-1/*")
948 .header(axum::http::header::IF_MATCH, etag)
949 .body(Body::from(Bytes::from_static(b"new\n")))
950 .unwrap(),
951 )
952 .await
953 .unwrap();
954
955 assert_eq!(response.status(), StatusCode::OK);
956 }
957
958 #[tokio::test]
959 async fn put_and_delete_on_replace_prefixed_path_still_work() {
960 let router = router();
961 let put = put_object(router.clone(), "replace/delete.md", b"gone").await;
962
963 assert_eq!(put.status(), StatusCode::CREATED);
964 let delete = router
965 .oneshot(
966 Request::builder()
967 .method("DELETE")
968 .uri(format!("/v1/knowledgebases/{KB}/replace/delete.md"))
969 .header("authorization", format!("Bearer {TOKEN}"))
970 .body(Body::empty())
971 .unwrap(),
972 )
973 .await
974 .unwrap();
975 assert_eq!(delete.status(), StatusCode::NO_CONTENT);
976 }
977
978 #[tokio::test]
979 async fn post_on_non_replace_path_returns_404_not_found() {
980 let response = router()
981 .oneshot(
982 Request::builder()
983 .method("POST")
984 .uri(format!("/v1/knowledgebases/{KB}/foo.md"))
985 .header("authorization", format!("Bearer {TOKEN}"))
986 .body(Body::empty())
987 .unwrap(),
988 )
989 .await
990 .unwrap();
991
992 assert_eq!(response.status(), StatusCode::NOT_FOUND);
993 let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
994 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
995 assert_eq!(json["error"], "not_found");
996 assert!(
997 json["message"]
998 .as_str()
999 .unwrap()
1000 .contains("supported actions: 'replace/<path>'")
1001 );
1002 }
1003
1004 #[tokio::test]
1005 async fn post_on_replace_prefixed_path_dispatches_to_replace_handler() {
1006 let router = router();
1007 let etag = put_object_etag(router.clone(), "target.md", b"hello world").await;
1008
1009 let response = post_replace(
1010 router,
1011 "target.md",
1012 &etag,
1013 br#"{"old_string":"world","new_string":"planet"}"#,
1014 )
1015 .await;
1016
1017 assert_eq!(response.status(), StatusCode::OK);
1018 let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
1019 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
1020 assert_eq!(json["match_count"], 1);
1021 }
1022
1023 #[tokio::test]
1024 async fn post_on_replace_replace_path_targets_the_replace_prefixed_object() {
1025 let router = router();
1026 let etag = put_object_etag(router.clone(), "replace/nested.md", b"foo bar").await;
1027
1028 let response = post_replace(
1029 router.clone(),
1030 "replace/nested.md",
1031 &etag,
1032 br#"{"old_string":"bar","new_string":"baz"}"#,
1033 )
1034 .await;
1035
1036 assert_eq!(response.status(), StatusCode::OK);
1037 let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
1038 let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
1039 assert_eq!(json["match_count"], 1);
1040 let get = get_object(router, "replace/nested.md").await;
1041 let body = to_bytes(get.into_body(), 64 * 1024).await.unwrap();
1042 assert_eq!(&body[..], b"foo baz");
1043 }
1044
1045 #[tokio::test]
1046 async fn test_conditional_put_503_then_naive_retry_412_keeps_object_stored() {
1047 let kb = KbSlug::try_new(KB).unwrap();
1048 let object_path = ObjectPath::try_from_str("cond.md").unwrap();
1049 let storage = Arc::new(crate::testing::InMemoryStorage::default());
1050
1051 let (indexer_tx, _rx) = tokio::sync::mpsc::channel(1);
1052 indexer_tx
1053 .try_send(IndexEvent::Upsert {
1054 kb: kb.clone(),
1055 object_key: ObjectPath::try_from_str("queued.md").unwrap(),
1056 etag: "etag".to_string(),
1057 mtime: 0,
1058 })
1059 .unwrap();
1060
1061 let mut kbs = BTreeMap::new();
1062 kbs.insert(KB.to_string(), kb.clone());
1063 let state = AppState {
1064 storage: storage.clone(),
1065 declared_kbs: Arc::new(kbs),
1066 public_read_policies: Arc::new(BTreeMap::new()),
1067 bearer_token: Arc::new(TOKEN.to_string()),
1068 max_body_size: MAX_BODY_BYTES,
1069 max_patchable_size: MAX_BODY_BYTES,
1070 indexer_tx,
1071 searcher: Arc::new(crate::testing::NoopSearcher),
1072 };
1073 let router = build_router(state);
1074
1075 let response = router
1076 .clone()
1077 .oneshot(
1078 Request::builder()
1079 .method("PUT")
1080 .uri(format!("/v1/knowledgebases/{KB}/cond.md"))
1081 .header("authorization", format!("Bearer {TOKEN}"))
1082 .header("if-none-match", "*")
1083 .body(Body::from(Bytes::from_static(b"first content")))
1084 .unwrap(),
1085 )
1086 .await
1087 .unwrap();
1088
1089 assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
1090 assert_eq!(response.headers().get("retry-after").unwrap(), "5");
1091 let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
1092 let body = String::from_utf8(body.to_vec()).unwrap();
1093 assert!(body.contains("\"error\":\"backend_unavailable\""));
1094 assert!(body.contains("object stored; indexer queue full — retry to re-enqueue"));
1095
1096 let stored = storage
1097 .get_object(&kb, &object_path, None, ConditionalHeaders::default())
1098 .await
1099 .unwrap();
1100 assert_eq!(stored.bytes, Bytes::from_static(b"first content"));
1101
1102 let retry = router
1103 .oneshot(
1104 Request::builder()
1105 .method("PUT")
1106 .uri(format!("/v1/knowledgebases/{KB}/cond.md"))
1107 .header("authorization", format!("Bearer {TOKEN}"))
1108 .header("if-none-match", "*")
1109 .body(Body::from(Bytes::from_static(b"second content")))
1110 .unwrap(),
1111 )
1112 .await
1113 .unwrap();
1114
1115 assert_eq!(retry.status(), StatusCode::PRECONDITION_FAILED);
1116 assert!(retry.headers().get("retry-after").is_none());
1117
1118 let stored = storage
1119 .get_object(&kb, &object_path, None, ConditionalHeaders::default())
1120 .await
1121 .unwrap();
1122 assert_eq!(stored.bytes, Bytes::from_static(b"first content"));
1123 }
1124
1125 #[tokio::test]
1126 async fn test_delete_returns_delete_specific_503_body_when_indexer_backpressure() {
1127 let kb = KbSlug::try_new(KB).unwrap();
1128 let object_path = ObjectPath::try_from_str("to-delete.md").unwrap();
1129 let storage = Arc::new(crate::testing::InMemoryStorage::default());
1130 storage
1131 .put_object(
1132 &kb,
1133 &object_path,
1134 Bytes::from_static(b"content"),
1135 Some("text/plain"),
1136 ConditionalHeaders::default(),
1137 )
1138 .await
1139 .unwrap();
1140
1141 let (indexer_tx, _rx) = tokio::sync::mpsc::channel(1);
1142 indexer_tx
1143 .try_send(IndexEvent::Upsert {
1144 kb: kb.clone(),
1145 object_key: ObjectPath::try_from_str("queued.md").unwrap(),
1146 etag: "etag".to_string(),
1147 mtime: 0,
1148 })
1149 .unwrap();
1150
1151 let mut kbs = BTreeMap::new();
1152 kbs.insert(KB.to_string(), kb.clone());
1153 let state = AppState {
1154 storage: storage.clone(),
1155 declared_kbs: Arc::new(kbs),
1156 public_read_policies: Arc::new(BTreeMap::new()),
1157 bearer_token: Arc::new(TOKEN.to_string()),
1158 max_body_size: MAX_BODY_BYTES,
1159 max_patchable_size: MAX_BODY_BYTES,
1160 indexer_tx,
1161 searcher: Arc::new(crate::testing::NoopSearcher),
1162 };
1163 let router = build_router(state);
1164
1165 let response = router
1166 .oneshot(
1167 Request::builder()
1168 .method("DELETE")
1169 .uri(format!("/v1/knowledgebases/{KB}/to-delete.md"))
1170 .header("authorization", format!("Bearer {TOKEN}"))
1171 .body(Body::empty())
1172 .unwrap(),
1173 )
1174 .await
1175 .unwrap();
1176
1177 assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
1178 assert_eq!(response.headers().get("retry-after").unwrap(), "5");
1179 let body = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
1180 let body = String::from_utf8(body.to_vec()).unwrap();
1181 assert!(body.contains("\"error\":\"backend_unavailable\""));
1182 assert!(
1183 body.contains("\"message\":\"deleted from storage; retry to clear from search index\"")
1184 );
1185 assert!(!body.contains("object stored; indexer queue full — retry to re-enqueue"));
1186
1187 let deleted = storage
1188 .get_object(&kb, &object_path, None, ConditionalHeaders::default())
1189 .await;
1190 assert!(matches!(deleted, Err(StorageError::NotFound { .. })));
1191 }
1192}