diff --git a/crates/didbot-pds/tests/blob_checkpoints.rs b/crates/didbot-pds/tests/blob_checkpoints.rs new file mode 100644 index 00000000..7c1ef6b6 --- /dev/null +++ b/crates/didbot-pds/tests/blob_checkpoints.rs @@ -0,0 +1,227 @@ +//! An upload judged while its body arrives. +//! +//! `blob.write` is asked about an upload before its body is read, at +//! checkpoints while the body arrives, and once it is in. A checkpoint is +//! there to refuse early: it stops an upload a policy refuses before the +//! rest of the body is read, it costs judgments in proportion to the log of +//! the body rather than to its frames, and it neither admits an upload nor +//! counts as the upload's attempt. + +use std::sync::{Arc, Mutex}; + +use didbot_identity::Zone; +use didbot_pds::policy::{Outcome, PolicyGate, PolicyVersion, Subject}; +use didbot_pds::{BlobError, MemoryAccountStore, ProvisionRequest, Provisioner, Registry}; + +const ZONE_HOST: &str = "agents.localhost"; +/// A PNG's own signature, which is what an upload is sniffed by. +const PNG: &[u8] = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR"; +const REASON: &str = "this blob is refused for its size"; + +/// One judgment, as the gate was asked it. +#[derive(Debug, Clone, PartialEq, Eq)] +struct Asked { + sniffed: Option, + size: Option, +} + +/// A gate refusing every upload whose size `refuses`, which writes down each +/// judgment it is asked and each attempt it observes. +struct SizeGate { + refuses: fn(u64) -> bool, + asked: Mutex>, + observed: Mutex>, +} + +impl SizeGate { + fn new(refuses: fn(u64) -> bool) -> Arc { + Arc::new(Self { + refuses, + asked: Mutex::default(), + observed: Mutex::default(), + }) + } + + fn asked(&self) -> Vec { + self.asked.lock().expect("poisoned").clone() + } + + fn observed(&self) -> Vec { + self.observed.lock().expect("poisoned").clone() + } +} + +impl PolicyGate for SizeGate { + fn judge(&self, subject: &Subject<'_>) -> Outcome { + let Subject::Blob { sniffed, size, .. } = subject else { + return Outcome::Allow; + }; + self.asked.lock().expect("poisoned").push(Asked { + sniffed: sniffed.map(str::to_owned), + size: *size, + }); + if size.is_some_and(self.refuses) { + Outcome::Reject { + reason: REASON.to_owned(), + } + } else { + Outcome::Allow + } + } + + fn version(&self) -> PolicyVersion { + PolicyVersion::none() + } + + fn observe(&self, subject: &Subject<'_>, outcome: &Outcome) { + if matches!(subject, Subject::Blob { .. }) { + self.observed + .lock() + .expect("poisoned") + .push(outcome.clone()); + } + } +} + +fn provisioner(gate: Arc) -> Provisioner { + Provisioner::new( + "did:web:operator.example", + Zone::delegated("localhost", ZONE_HOST).expect("a valid zone"), + "http://localhost:3600".to_string(), + MemoryAccountStore::new(), + ) + .with_policy_gate(gate) +} + +fn account(pds: &impl Registry) -> String { + pds.provision(ProvisionRequest::new("kestrel", None)) + .expect("provisioning succeeds") + .account + .did + .as_str() + .to_owned() +} + +/// `len` bytes that open with a PNG's signature. +fn png(len: usize) -> Vec { + let mut body = vec![0u8; len]; + body[..PNG.len()].copy_from_slice(PNG); + body +} + +/// Sixteen times the body in sixteen times the frames costs four more +/// judgments, every one after the first has the type the bytes carry, and +/// the upload is observed once. +#[test] +fn judgments_grow_with_the_log_of_the_body_and_not_its_frames() { + let judged = |len: usize| { + let gate = SizeGate::new(|_| false); + let pds = provisioner(gate.clone()); + let did = account(&pds); + let mut upload = pds + .begin_blob(&did, "image/png", None) + .expect("the upload opens"); + for frame in png(len).chunks(1024) { + upload.write(frame).expect("a frame writes"); + } + let reference = upload.commit().expect("the upload commits"); + assert_eq!(reference.size, len as u64); + assert_eq!( + gate.observed(), + vec![Outcome::Allow], + "one upload is one attempt" + ); + gate.asked() + }; + + let small = judged(256 * 1024); + let large = judged(4 * 1024 * 1024); + assert_eq!(large.len() - small.len(), 4, "{small:?}\n{large:?}"); + assert_eq!( + large[0], + Asked { + sniffed: None, + size: None + }, + "the first judgment is before the body" + ); + assert!( + large[1..] + .iter() + .all(|asked| asked.sniffed.as_deref() == Some("image/png")), + "{large:?}" + ); +} + +/// A body a policy refuses for its size is refused within twice the limit, +/// with no length declared to refuse it by, and the refusal is the one +/// attempt observed. Nothing is stored. +#[test] +fn a_checkpoint_refuses_an_upload_once_it_passes_the_limit() { + const LIMIT: u64 = 100_000; + let gate = SizeGate::new(|size| size > LIMIT); + let pds = provisioner(gate.clone()); + let did = account(&pds); + let held = || { + pds.list_blobs(&did, 10, None, None) + .expect("the account lists its blobs") + .0 + }; + let before = held(); + let mut upload = pds + .begin_blob(&did, "image/png", None) + .expect("no length is declared, so nothing refuses it yet"); + + let mut written = 0; + let refused = png(4 * 1024 * 1024).chunks(16 * 1024).find_map(|frame| { + written += frame.len() as u64; + upload.write(frame).err() + }); + assert_eq!( + refused, + Some(BlobError::PolicyRejected { + reason: REASON.to_owned() + }) + ); + assert!( + written <= 2 * LIMIT, + "{written} bytes were read of a body refused past {LIMIT}" + ); + assert_eq!( + gate.observed(), + vec![Outcome::Reject { + reason: REASON.to_owned() + }], + "one upload is one attempt" + ); + + drop(upload); + assert_eq!(held(), before, "the refused upload left a blob behind"); +} + +/// An upload that declared its length is judged at each checkpoint as that +/// length, which is what it comes to once it is in. A rule refusing small +/// blobs does not refuse a large one part of the way through it. +#[test] +fn a_checkpoint_judges_the_length_an_upload_declared() { + const LEN: usize = 300 * 1024; + let gate = SizeGate::new(|size| size < 200 * 1024); + let pds = provisioner(gate.clone()); + let did = account(&pds); + let mut upload = pds + .begin_blob(&did, "image/png", Some(LEN as u64)) + .expect("the declared length is allowed"); + for frame in png(LEN).chunks(16 * 1024) { + upload + .write(frame) + .expect("no checkpoint refuses the declared length"); + } + let reference = upload.commit().expect("the upload commits"); + assert_eq!(reference.size, LEN as u64); + let asked = gate.asked(); + assert!(asked.len() > 2, "no checkpoint was judged: {asked:?}"); + assert!( + asked.iter().all(|asked| asked.size == Some(LEN as u64)), + "{asked:?}" + ); +} diff --git a/crates/didbot-serve/tests/blob_policy.rs b/crates/didbot-serve/tests/blob_policy.rs index e13c29e7..73391908 100644 --- a/crates/didbot-serve/tests/blob_policy.rs +++ b/crates/didbot-serve/tests/blob_policy.rs @@ -1,12 +1,18 @@ -//! A policy on `blob.write` judges an upload by what its bytes are. +//! A policy on `blob.write` judges an upload by what its bytes are, and by +//! its size while they arrive. mod support; use std::collections::BTreeMap; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; +use std::task::{Context, Poll}; +use axum::body::Body; use axum::http::StatusCode; -use serde_json::json; +use futures_core::Stream; +use serde_json::{json, Value}; use didbot_serve::oauth::authorize::GrantAnyScope; use support::{Harness, OPERATOR}; @@ -14,23 +20,19 @@ use support::{Harness, OPERATOR}; const REASON: &str = "only images may be uploaded"; /// A PNG's own signature, which is what an upload is sniffed by. const PNG: &[u8] = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR"; +/// How long each frame of a streamed body is. +const FRAME: usize = 16 * 1024; -/// A deployment refusing every upload whose bytes are not an image. -fn images_only() -> Harness { +/// A deployment whose one policy on `blob.write` is `statement`. +fn blob_policy(statement: Value) -> Harness { let harness = Harness::build(&[], Arc::new(GrantAnyScope)); let policies = BTreeMap::from([( - "images-only".to_owned(), + "blobs".to_owned(), json!({ "actions": ["blob.write"], "document": { "$type": "bot.did.policy#regex", - "statements": [{ - "kind": "denyBlobUnless", - "mimeType": [{"regex": "image/.*"}], - "sniffed": [{"regex": "image/.*"}], - "maxSize": 1000, - "reason": REASON, - }], + "statements": [statement], }, "createdAt": "2026-09-11T10:00:00Z", }), @@ -38,7 +40,7 @@ fn images_only() -> Harness { let bindings = BTreeMap::from([( "everyone".to_owned(), json!({ - "policies": [format!("at://{OPERATOR}/bot.did.policy/images-only")], + "policies": [format!("at://{OPERATOR}/bot.did.policy/blobs")], "subjects": [harness.registry.service_did()], "includes": ["descendants"], "createdAt": "2026-09-11T10:00:00Z", @@ -48,6 +50,65 @@ fn images_only() -> Harness { harness } +/// A deployment refusing every upload whose bytes are not an image, or that +/// is over a thousand bytes. +fn images_only() -> Harness { + blob_policy(json!({ + "kind": "denyBlobUnless", + "mimeType": [{"regex": "image/.*"}], + "sniffed": [{"regex": "image/.*"}], + "maxSize": 1000, + "reason": REASON, + })) +} + +/// `len` bytes that open with a PNG's signature. +fn png(len: usize) -> Vec { + let mut bytes = vec![0u8; len]; + bytes[..PNG.len()].copy_from_slice(PNG); + bytes +} + +/// A body handed over a [`FRAME`] at a time, counting the bytes the server +/// has taken from it. +struct Counted { + frames: std::vec::IntoIter>, + taken: Arc, +} + +impl Stream for Counted { + type Item = Result, std::io::Error>; + + fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + let frame = this.frames.next(); + if let Some(frame) = &frame { + this.taken.fetch_add(frame.len(), Ordering::SeqCst); + } + Poll::Ready(frame.map(Ok)) + } +} + +/// `bytes` as a streamed body, and the count of what the server takes. +fn streamed(bytes: &[u8]) -> (Body, Arc) { + let taken = Arc::new(AtomicUsize::new(0)); + let frames: Vec> = bytes.chunks(FRAME).map(<[u8]>::to_vec).collect(); + let body = Body::from_stream(Counted { + frames: frames.into_iter(), + taken: taken.clone(), + }); + (body, taken) +} + +/// The blobs the harness's account holds. +fn held(harness: &Harness) -> Vec { + harness + .registry + .list_blobs(&harness.account_did, 100, None, None) + .expect("the account lists its blobs") + .0 +} + /// A type or a length the request itself declares is judged before the body /// is read, so a blob a policy will refuse is never carried. #[tokio::test] @@ -100,3 +161,90 @@ async fn an_upload_nothing_judges_is_stored() { let (status, body) = harness.upload("image/png", b"not an image".to_vec()).await; assert_eq!(status, StatusCode::OK, "{body}"); } + +/// A body a size policy refuses is refused while it streams, even with no +/// `Content-Length` to judge it by: by the first checkpoint, well short of +/// the server's own blob limit. +#[tokio::test] +async fn a_size_policy_refuses_an_upload_with_no_length_before_it_is_read() { + let harness = images_only(); + let before = held(&harness); + let body = png(4 * 1024 * 1024); + let (stream, taken) = streamed(&body); + + let (status, answer) = harness.upload_body("image/png", None, stream).await; + assert_eq!(status, StatusCode::FORBIDDEN, "{answer}"); + assert_eq!(answer["error"], "PolicyViolation"); + assert!( + answer["message"] + .as_str() + .is_some_and(|m| m.contains(REASON)), + "{answer}" + ); + let taken = taken.load(Ordering::SeqCst); + assert!( + taken <= 64 * 1024 + FRAME, + "{taken} of {} bytes were read before the refusal", + body.len() + ); + assert_eq!( + harness + .denials() + .iter() + .filter(|row| row.reason == REASON) + .count(), + 1, + "the refusal is written down once" + ); + assert_eq!(held(&harness), before, "the refused upload left a blob"); +} + +/// A body that is not the length its `Content-Length` declares is refused +/// with no policy loaded: one that runs past it at the frame that does, and +/// one that ends short before it is stored. +#[tokio::test] +async fn a_body_that_is_not_its_content_length_is_refused() { + let harness = Harness::build(&[], Arc::new(GrantAnyScope)); + let before = held(&harness); + + let (stream, taken) = streamed(&png(4 * 1024 * 1024)); + let (status, answer) = harness.upload_body("image/png", Some(100), stream).await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{answer}"); + assert_eq!(answer["error"], "InvalidRequest"); + assert_eq!( + taken.load(Ordering::SeqCst), + FRAME, + "nothing past the frame that ran past the length is read" + ); + + let (stream, _) = streamed(&png(100)); + let (status, answer) = harness.upload_body("image/png", Some(1000), stream).await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{answer}"); + assert_eq!(answer["error"], "InvalidRequest"); + + assert_eq!(held(&harness), before, "a refused upload left a blob"); +} + +/// An upload under a size policy's limit passes every checkpoint and is +/// stored as it was sent, with or without a `Content-Length`. +#[tokio::test] +async fn an_upload_under_the_limit_is_stored_whole() { + let harness = blob_policy(json!({ + "kind": "denyBlobUnless", + "maxSize": 1024 * 1024, + "reason": REASON, + })); + let body = png(300 * 1024); + for declared in [None, Some(body.len() as u64)] { + let (stream, _) = streamed(&body); + let (status, answer) = harness.upload_body("image/png", declared, stream).await; + assert_eq!(status, StatusCode::OK, "{answer}"); + assert_eq!(answer["blob"]["size"], body.len()); + assert_eq!(answer["blob"]["mimeType"], "image/png"); + assert_eq!( + answer["blob"]["ref"]["$link"], + didbot_data::Cid::of_raw(&body).to_string() + ); + } + assert!(harness.denials().is_empty(), "{:?}", harness.denials()); +} diff --git a/crates/didbot-serve/tests/support/mod.rs b/crates/didbot-serve/tests/support/mod.rs index 7fed3079..d4570eef 100644 --- a/crates/didbot-serve/tests/support/mod.rs +++ b/crates/didbot-serve/tests/support/mod.rs @@ -454,6 +454,26 @@ impl Harness { .await } + /// An upload of `body` as `content_type` with the agent's own token, + /// declaring `content_length` when there is one. + pub async fn upload_body( + &self, + content_type: &str, + content_length: Option, + body: Body, + ) -> (StatusCode, Value) { + let mut request = Request::builder() + .method("POST") + .uri("/xrpc/com.atproto.repo.uploadBlob") + .header("authorization", self.as_agent()) + .header("dpop", "any-proof") + .header("content-type", content_type); + if let Some(length) = content_length { + request = request.header("content-length", length); + } + self.call(request.body(body).unwrap()).await + } + /// A `createRecord` of [`thing`] into the agent's repository. pub async fn create(&self, authorization: &str) -> (StatusCode, Value) { self.xrpc_post(