use std::collections::BTreeSet;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime};
use http::StatusCode;
use knot_record::attachment::{
Attachment, AttachmentMime, AttachmentSize, AttachmentTooBig, MAX_UPLOAD_BYTES, OCTET_STREAM,
SNIFF_BYTES, sniffed,
};
use knot_resource::{BelowFloor, DiskFloorBytes, DiskGovernor, ReserveBytes, disk_free_bytes};
use knot_types::{BlobCid, RepoDid, Sha256Digest};
use sha2::{Digest as _, Sha256};
use tempfile::NamedTempFile;
use crate::cob_write::WriteError;
use crate::datastore::marks::expire;
use crate::plan::repo_key;
use crate::repolock::{self, LockMode};
pub use crate::xrpc::com_atproto_sync_get_blob::{emit, locations};
pub const UPLOAD_BLOB: &str = "com.atproto.repo.uploadBlob";
pub(crate) const BLOBS_DIR: &str = "www/_blobs";
const UNREFERENCED_GRACE: Duration = Duration::from_secs(14 * 24 * 60 * 60);
const SPOOL_LIFETIME: Duration = Duration::from_secs(24 * 60 * 60);
const SPOOL_DIR: &str = ".incoming";
const FREE_SPACE_FLOOR: u64 = 1_073_741_824;
const FREE_SPACE_FLOOR_ENV: &str = "KNOT_LFS_FREE_SPACE_FLOOR_BYTES";
pub(crate) const SERVED: [(&str, &str); 9] = [
("image/png", "png"),
("image/jpeg", "jpg"),
("image/gif", "gif"),
("image/webp", "webp"),
("image/bmp", "bmp"),
("application/pdf", "pdf"),
("video/mp4", "mp4"),
("image/svg+xml", "svg"),
(OCTET_STREAM, "bin"),
];
pub struct Store {
dir: PathBuf,
}
impl Store {
pub fn of(output_root: &Path, repo: &RepoDid) -> Self {
Self {
dir: repolock::repo_dir(output_root, &repo_key(repo.as_str())).join("blobs"),
}
}
pub(crate) fn path(&self, cid: &BlobCid) -> PathBuf {
self.dir.join(cid.to_string())
}
fn spool(&self) -> std::io::Result<NamedTempFile> {
let spool = self.dir.join(SPOOL_DIR);
std::fs::create_dir_all(&spool)?;
expire(&spool, SPOOL_LIFETIME);
tempfile::Builder::new().tempfile_in(spool)
}
pub fn size(&self, cid: &BlobCid) -> std::io::Result<Option<u64>> {
match std::fs::metadata(self.path(cid)) {
Ok(meta) => Ok(Some(meta.len())),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(error) => Err(error),
}
}
pub fn sniffed(&self, cid: &BlobCid) -> Option<&'static str> {
let mut prefix = Vec::new();
std::fs::File::open(self.path(cid))
.ok()?
.take(SNIFF_BYTES as u64)
.read_to_end(&mut prefix)
.ok()?;
sniffed(&prefix)
}
pub fn touch(&self, cid: &BlobCid) -> std::io::Result<()> {
std::fs::File::options()
.write(true)
.open(self.path(cid))?
.set_modified(SystemTime::now())
}
pub(crate) fn sweep(
&self,
referenced: &BTreeSet<String>,
now: SystemTime,
) -> std::io::Result<()> {
let entries = match std::fs::read_dir(&self.dir) {
Ok(entries) => entries,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(error) => return Err(error),
};
for entry in entries.flatten() {
let name = entry.file_name();
let untouched = entry
.metadata()
.and_then(|meta| meta.modified())
.ok()
.and_then(|modified| now.duration_since(modified).ok())
.is_some_and(|age| age >= UNREFERENCED_GRACE);
let referenced = name.to_str().is_none_or(|name| referenced.contains(name));
if entry.path().is_file() && !referenced && untouched {
match std::fs::remove_file(entry.path()) {
Err(error) if error.kind() != std::io::ErrorKind::NotFound => {
return Err(error);
}
_ => {}
}
}
}
Ok(())
}
}
fn room_for(dir: &Path, bytes: u64) -> Result<(), WriteError> {
let floor = std::env::var(FREE_SPACE_FLOOR_ENV)
.ok()
.and_then(|floor| floor.parse().ok())
.unwrap_or(FREE_SPACE_FLOOR);
let free = disk_free_bytes(dir).map_err(|error| WriteError::internal(error.to_string()))?;
DiskGovernor::new(DiskFloorBytes::new(floor))
.reserve_against(free, ReserveBytes::new(bytes))
.map(|_| ())
.map_err(|BelowFloor { free, floor }| {
WriteError::unsupported(
StatusCode::SERVICE_UNAVAILABLE,
"BlobStoreFull",
format!(
"The blob store's disk is down to {} bytes, under the {}-byte floor. \
Please try again once there's room, thanks.",
free.get(),
floor.get()
),
)
})
}
struct Received {
size: u64,
digest: Sha256Digest,
prefix: Vec<u8>,
}
fn receive(
body: impl Read,
limit: u64,
spool: &mut impl std::io::Write,
) -> std::io::Result<Received> {
let mut hasher = Sha256::new();
let mut prefix = Vec::with_capacity(SNIFF_BYTES);
let mut size = 0;
let mut body = body.take(limit + 1);
let mut chunk = [0u8; 64 * 1024];
loop {
let read = body.read(&mut chunk)?;
if read == 0 {
break;
}
hasher.update(&chunk[..read]);
prefix.extend_from_slice(&chunk[..read.min(SNIFF_BYTES - prefix.len())]);
spool.write_all(&chunk[..read])?;
size += read as u64;
}
Ok(Received {
size,
digest: Sha256Digest::new(hasher.finalize().into()),
prefix,
})
}
pub fn upload(
output_root: &Path,
repo: &RepoDid,
body: impl Read,
declared_length: Option<usize>,
content_type: &str,
) -> Result<Attachment, WriteError> {
let too_large = |sent: u64| {
WriteError::request_too_large(format!(
"An attachment is at most {MAX_UPLOAD_BYTES} bytes, and this upload sends {sent}."
))
};
if let Some(length) = declared_length.filter(|length| *length as u64 > MAX_UPLOAD_BYTES) {
return Err(too_large(length as u64));
}
let internal = |error: std::io::Error| WriteError::internal(error.to_string());
let store = Store::of(output_root, repo);
let mut spool = store.spool().map_err(internal)?;
room_for(
&store.dir,
declared_length.map_or(0, |length| length as u64),
)?;
let received = receive(body, MAX_UPLOAD_BYTES, &mut spool).map_err(internal)?;
if let Some(length) = declared_length.filter(|length| *length as u64 != received.size) {
return Err(WriteError::invalid_request(format!(
"This upload's body is {} bytes, and its Content-Length says {length}. \
Please send the whole body.",
received.size
)));
}
if received.size == 0 {
return Err(WriteError::invalid_request(
"This upload's body is empty; please send the attachment bytes as the body.",
));
}
let size = AttachmentSize::new(received.size)
.map_err(|AttachmentTooBig { sent, .. }| too_large(sent))?;
let cid = BlobCid::from_digest(received.digest);
spool.as_file().sync_all().map_err(internal)?;
// the spool holds the bytes already, and only the free-space floor is checked.
room_for(&store.dir, 0)?;
// a sweep that judged this blob stale can't unlink the copy being stored.
let _lock = repolock::acquire(output_root, &repo_key(repo.as_str()), LockMode::Exclusive)
.map_err(|error| WriteError::internal(format!("{error:#}")))?;
spool
.persist(store.path(&cid))
.map_err(|error| internal(error.error))?;
Ok(Attachment::new(
cid,
AttachmentMime::detected(&received.prefix, content_type),
size,
))
}
fn references(value: &serde_json::Value, into: &mut BTreeSet<String>) {
match value {
serde_json::Value::Object(fields) => {
if fields.get("$type").is_some_and(|kind| kind == "blob")
&& let Some(link) = fields["ref"]["$link"].as_str()
{
into.insert(link.to_owned());
}
fields.values().for_each(|field| references(field, into));
}
serde_json::Value::Array(items) => items.iter().for_each(|item| references(item, into)),
_ => {}
}
}
pub(crate) fn referenced<'a>(
values: impl IntoIterator<Item = &'a serde_json::Value>,
) -> BTreeSet<String> {
let mut into = BTreeSet::new();
values
.into_iter()
.for_each(|value| references(value, &mut into));
into
}
#[cfg(test)]
mod tests {
use super::*;
fn cid_of(bytes: &[u8]) -> BlobCid {
BlobCid::from_digest(Sha256Digest::hash(bytes))
}
fn store_with(blobs: &[&[u8]]) -> (tempfile::TempDir, Store) {
let dir = tempfile::tempdir().unwrap();
let store = Store {
dir: dir.path().to_path_buf(),
};
for bytes in blobs {
std::fs::write(store.path(&cid_of(bytes)), bytes).unwrap();
}
(dir, store)
}
fn age(store: &Store, bytes: &[u8], by: Duration) {
std::fs::File::options()
.write(true)
.open(store.path(&cid_of(bytes)))
.unwrap()
.set_modified(SystemTime::now() - by)
.unwrap();
}
#[test]
fn every_type_knot_record_sniffs_has_a_served_extension() {
let samples: [&[u8]; 8] = [
b"\x89PNG\r\n",
b"\xFF\xD8\xFF\xE0",
b"GIF89a",
b"RIFF\0\0\0\0WEBPVP8 ",
b"BM\x01\0\0\0",
b"%PDF-1.7",
b"\0\0\0\x18ftypmp42",
b"<svg xmlns=\"x\"/>",
];
for sample in samples {
let mime = sniffed(sample).expect("a sniffed sample");
assert!(
SERVED.iter().any(|(served, _)| *served == mime),
"{mime} has no extension for nginx to type it by"
);
}
assert!(SERVED.iter().any(|(mime, _)| *mime == OCTET_STREAM));
}
#[test]
fn a_record_references_the_blobs_in_every_shape_it_carries_them() {
let blob = |link: &str| serde_json::json!({ "$type": "blob", "ref": { "$link": link }, "mimeType": "image/png", "size": 1 });
let record = serde_json::json!({
"artifact": blob("artifact"),
"body": { "text": "x", "blobs": [{ "$type": "embed", "image": blob("inline") }] },
"blobs": [{ "$type": "embed", "video": blob("attached") }],
"subject": { "$type": "org.tangled.git.oid", "ref": { "$link": "not-a-blob" } },
});
let mut found = BTreeSet::new();
references(&record, &mut found);
assert_eq!(
found,
BTreeSet::from(["artifact".into(), "inline".into(), "attached".into()])
);
}
#[test]
fn receive_reads_one_byte_past_the_limit_and_no_more() {
let mut spooled = Vec::new();
let received = receive(std::io::repeat(7), 10, &mut spooled).unwrap();
assert_eq!((received.size, spooled.len()), (11, 11));
let small = receive(&b"abc"[..], 10, &mut Vec::new()).unwrap();
assert_eq!(small.size, 3);
assert_eq!(small.digest, Sha256Digest::hash(b"abc"));
}
#[test]
fn a_sweep_takes_only_unreferenced_blobs_past_the_grace_window() {
let (_dir, store) = store_with(&[b"stale", b"young", b"cited"]);
age(
&store,
b"stale",
UNREFERENCED_GRACE + Duration::from_secs(60),
);
age(
&store,
b"cited",
UNREFERENCED_GRACE + Duration::from_secs(60),
);
age(
&store,
b"young",
UNREFERENCED_GRACE - Duration::from_secs(60),
);
let referenced = BTreeSet::from([cid_of(b"cited").to_string()]);
store.sweep(&referenced, SystemTime::now()).unwrap();
assert_eq!(store.size(&cid_of(b"stale")).unwrap(), None);
assert_eq!(store.size(&cid_of(b"young")).unwrap(), Some(5));
assert_eq!(store.size(&cid_of(b"cited")).unwrap(), Some(5));
}
#[test]
fn touching_a_blob_restarts_its_grace_window() {
let (_dir, store) = store_with(&[b"adopted"]);
age(&store, b"adopted", UNREFERENCED_GRACE * 2);
store.touch(&cid_of(b"adopted")).unwrap();
store.sweep(&BTreeSet::new(), SystemTime::now()).unwrap();
assert_eq!(store.size(&cid_of(b"adopted")).unwrap(), Some(7));
}
#[test]
fn a_new_upload_clears_spools_left_by_dead_ones() {
let (_dir, store) = store_with(&[]);
let spool = store.dir.join(SPOOL_DIR);
std::fs::create_dir_all(&spool).unwrap();
let (dead, live) = (spool.join("dead"), spool.join("live"));
std::fs::write(&dead, b"half").unwrap();
std::fs::write(&live, b"going").unwrap();
std::fs::File::options()
.write(true)
.open(&dead)
.unwrap()
.set_modified(SystemTime::now() - SPOOL_LIFETIME * 2)
.unwrap();
let _new = store.spool().unwrap();
assert!(!dead.exists() && live.exists());
}
}