1mod backend;
2mod codec;
3mod dedupe;
4mod rewrite;
5pub mod s3;
6mod staging;
7mod sweep;
8mod verify;
9mod walk;
10
11#[cfg(test)]
12mod tests;
13
14use std::collections::HashMap;
15use std::path::{Path, PathBuf};
16use std::sync::atomic::{AtomicU64, Ordering};
17use std::time::Instant;
18
19use tokio::sync::Mutex;
20
21use futures_util::{Stream, StreamExt};
22use sha2::{Digest, Sha256};
23use tokio::fs;
24use tokio::io::AsyncWriteExt;
25
26use crate::error::Error;
27use crate::namespace::Namespace;
28
29pub use backend::Store;
30pub use dedupe::DedupeReport;
31use dedupe::shares_bytes_with;
32pub use rewrite::CompressReport;
33pub use verify::VerifyReport;
34
35enum Sink {
36 Raw(fs::File),
37 Framed(Box<codec::Writer>),
38}
39
40impl Sink {
41 async fn write(&mut self, chunk: &[u8]) -> Result<(), Error> {
42 match self {
43 Self::Raw(file) => Ok(file.write_all(chunk).await?),
44 Self::Framed(writer) => writer.push(chunk).await,
45 }
46 }
47
48 async fn finish(self) -> Result<(), Error> {
49 match self {
50 Self::Raw(mut file) => {
51 file.flush().await?;
52 Ok(file.sync_all().await?)
53 }
54 Self::Framed(writer) => writer.finish().await,
55 }
56 }
57}
58
59pub struct Staged {
61 pub path: PathBuf,
62 destination: PathBuf,
63 pub written: u64,
64 fresh: bool,
65}
66
67pub enum Object {
69 Raw {
70 file: fs::File,
71 size: u64,
72 },
73 Framed(codec::Framed),
74 Remote {
75 bucket: s3::S3Store,
76 oid: String,
77 size: u64,
78 },
79}
80
81impl Object {
82 pub fn size(&self) -> u64 {
83 match self {
84 Self::Raw { size, .. } => *size,
85 Self::Framed(framed) => framed.plaintext(),
86 Self::Remote { size, .. } => *size,
87 }
88 }
89
90 pub async fn stream(
91 self,
92 start: u64,
93 length: u64,
94 ) -> Result<futures_util::stream::BoxStream<'static, Result<axum::body::Bytes, Error>>, Error>
95 {
96 use futures_util::StreamExt;
97 use tokio::io::AsyncSeekExt;
98
99 match self {
100 Self::Raw { mut file, .. } => {
101 file.seek(std::io::SeekFrom::Start(start)).await?;
102 let reader =
103 tokio_util::io::ReaderStream::new(tokio::io::AsyncReadExt::take(file, length));
104
105 Ok(reader.map(|chunk| chunk.map_err(Error::from)).boxed())
106 }
107 Self::Framed(framed) => Ok(framed.stream(start, length).boxed()),
108 Self::Remote { bucket, oid, .. } => {
109 let chunks = bucket.read(&oid, start, length).await?;
110
111 Ok(chunks
112 .map(|chunk| {
113 chunk.map_err(|error| Error::Storage(std::io::Error::other(error)))
114 })
115 .boxed())
116 }
117 }
118 }
119}
120pub use staging::{Reclaimed, reclaim};
121pub use sweep::SweepReport;
122
123#[derive(Debug, Clone, Copy)]
128pub struct Budget {
129 pub used: u64,
130 pub limit: u64,
131}
132
133impl Budget {
134 pub fn exceeded_by(&self, arriving: u64) -> bool {
135 self.used + arriving > self.limit
136 }
137
138 pub fn refusal(&self) -> Error {
139 Error::OverQuota {
140 used: self.used,
141 limit: self.limit,
142 }
143 }
144}
145
146pub struct LocalStore {
147 root: PathBuf,
148 counter: AtomicU64,
149 usage: Mutex<Option<(Instant, u64, u64)>>,
150 per_namespace: Mutex<HashMap<String, (Instant, u64, u64)>>,
151 scans: AtomicU64,
152 max_object_size: Option<u64>,
153 compression: Option<i32>,
154}
155
156impl LocalStore {
157 pub fn new(root: impl Into<PathBuf>) -> Self {
158 Self {
159 root: root.into(),
160 counter: AtomicU64::new(0),
161 usage: Mutex::new(None),
162 per_namespace: Mutex::new(HashMap::new()),
163 scans: AtomicU64::new(0),
164 max_object_size: None,
165 compression: None,
166 }
167 }
168
169 pub fn with_compression(mut self, level: Option<i32>) -> Self {
170 self.compression = level;
171 self
172 }
173
174 pub fn with_max_object_size(mut self, limit: Option<u64>) -> Self {
175 self.max_object_size = limit;
176 self
177 }
178
179 pub fn validate_oid(oid: &str) -> Result<(), Error> {
180 let well_formed = oid.len() == 64
181 && oid
182 .bytes()
183 .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b));
184
185 well_formed.then_some(()).ok_or(Error::MalformedOid)
186 }
187
188 fn object_path(&self, ns: &Namespace, oid: &str) -> PathBuf {
189 self.root
190 .join(ns.org())
191 .join(ns.repo())
192 .join(&oid[0..2])
193 .join(&oid[2..4])
194 .join(oid)
195 }
196
197 fn content_path(&self, oid: &str) -> PathBuf {
198 self.root
199 .join(".content")
200 .join(&oid[0..2])
201 .join(&oid[2..4])
202 .join(oid)
203 }
204
205 pub fn scans(&self) -> u64 {
206 self.scans.load(Ordering::Relaxed)
207 }
208
209 pub async fn writable(&self) -> Result<(), Error> {
210 fs::create_dir_all(&self.root).await?;
211
212 let ticket = self.counter.fetch_add(1, Ordering::Relaxed);
213 let probe = self.root.join(format!(".readiness.{ticket}"));
214
215 fs::write(&probe, b"").await?;
216 fs::remove_file(&probe).await?;
217
218 Ok(())
219 }
220
221 pub async fn exists(&self, ns: &Namespace, oid: &str) -> bool {
222 Self::validate_oid(oid).is_ok() && fs::metadata(self.object_path(ns, oid)).await.is_ok()
223 }
224
225 pub async fn open(&self, ns: &Namespace, oid: &str) -> Result<Object, Error> {
226 Self::validate_oid(oid)?;
227 let path = self.object_path(ns, oid);
228 let file = fs::File::open(&path).await.map_err(|_| Error::NotFound)?;
229 let on_disk = file.metadata().await?.len();
230
231 match codec::Framed::open(file, on_disk).await? {
232 Some(framed) => Ok(Object::Framed(framed)),
233 None => Ok(Object::Raw {
234 file: fs::File::open(&path).await.map_err(|_| Error::NotFound)?,
235 size: on_disk,
236 }),
237 }
238 }
239
240 pub async fn stage<S, E>(
246 &self,
247 ns: &Namespace,
248 oid: &str,
249 expected_size: Option<u64>,
250 budget: Option<Budget>,
251 mut chunks: S,
252 ) -> Result<Staged, Error>
253 where
254 S: Stream<Item = Result<axum::body::Bytes, E>> + Unpin,
255 E: std::error::Error + Send + Sync + 'static,
256 {
257 Self::validate_oid(oid)?;
258
259 if let Some(limit) = self.max_object_size
260 && expected_size.is_some_and(|declared| declared > limit)
261 {
262 return Err(Error::TooLarge { limit });
263 }
264
265 let path = self.object_path(ns, oid);
266 let parent = path.parent().expect("object paths always have a parent");
267 fs::create_dir_all(parent).await?;
268
269 let fresh = fs::metadata(&path).await.is_err();
272 let staged = self.staging_path(parent, oid);
273
274 match self.stream_to(&staged, budget, &mut chunks).await {
275 Ok((digest, written)) => {
276 if let Err(error) = Self::agrees(oid, expected_size, &digest, written) {
277 let _ = fs::remove_file(&staged).await;
278 return Err(error);
279 }
280
281 Ok(Staged {
282 path: staged,
283 destination: path,
284 written,
285 fresh,
286 })
287 }
288 Err(error) => {
289 let _ = fs::remove_file(&staged).await;
290 Err(error)
291 }
292 }
293 }
294
295 fn agrees(
296 oid: &str,
297 expected_size: Option<u64>,
298 digest: &str,
299 written: u64,
300 ) -> Result<(), Error> {
301 if let Some(declared) = expected_size.filter(|declared| *declared != written) {
302 return Err(Error::SizeMismatch {
303 declared,
304 actual: written,
305 });
306 }
307
308 if digest != oid {
309 return Err(Error::OidMismatch {
310 declared: oid.to_owned(),
311 actual: digest.to_owned(),
312 });
313 }
314
315 Ok(())
316 }
317
318 pub async fn write<S, E>(
319 &self,
320 ns: &Namespace,
321 oid: &str,
322 expected_size: Option<u64>,
323 budget: Option<Budget>,
324 chunks: S,
325 ) -> Result<u64, Error>
326 where
327 S: Stream<Item = Result<axum::body::Bytes, E>> + Unpin,
328 E: std::error::Error + Send + Sync + 'static,
329 {
330 let staged = self.stage(ns, oid, expected_size, budget, chunks).await?;
331
332 self.link_or_move(&staged.path, &staged.destination, oid)
333 .await?;
334
335 if staged.fresh {
336 self.stored(ns, staged.written).await;
337 }
338
339 Ok(staged.written)
340 }
341
342 fn staging_path(&self, parent: &Path, oid: &str) -> PathBuf {
343 let ticket = self.counter.fetch_add(1, Ordering::Relaxed);
344 parent.join(format!("{oid}.{ticket}.part"))
345 }
346
347 async fn stream_to<S, E>(
348 &self,
349 staged: &Path,
350 budget: Option<Budget>,
351 chunks: &mut S,
352 ) -> Result<(String, u64), Error>
353 where
354 S: Stream<Item = Result<axum::body::Bytes, E>> + Unpin,
355 E: std::error::Error + Send + Sync + 'static,
356 {
357 let file = fs::File::create(staged).await?;
358 let mut sink = match self.compression {
362 Some(level) => Sink::Framed(Box::new(codec::Writer::open(file, level).await?)),
363 None => Sink::Raw(file),
364 };
365 let mut hasher = Sha256::new();
366 let mut written = 0u64;
367
368 while let Some(chunk) = chunks.next().await {
369 let chunk = chunk.map_err(std::io::Error::other)?;
370 hasher.update(&chunk);
371 written += chunk.len() as u64;
372
373 if let Some(limit) = self.max_object_size.filter(|limit| written > *limit) {
378 return Err(Error::TooLarge { limit });
379 }
380
381 if let Some(budget) = budget.filter(|budget| budget.exceeded_by(written)) {
382 return Err(budget.refusal());
383 }
384
385 sink.write(&chunk).await?;
386 }
387
388 sink.finish().await?;
389
390 Ok((hex::encode(hasher.finalize()), written))
391 }
392
393 async fn link_or_move(&self, staged: &Path, final_path: &Path, oid: &str) -> Result<(), Error> {
400 let content = self.content_path(oid);
401 let parent = content.parent().expect("content paths have a parent");
402 fs::create_dir_all(parent).await?;
403
404 if fs::metadata(&content).await.is_err() {
405 fs::rename(staged, &content).await?;
406 }
407
408 match self.link(&content, final_path).await {
409 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
415 fs::rename(staged, &content).await?;
416 self.link(&content, final_path).await?;
417 }
418 outcome => outcome?,
419 }
420
421 let _ = fs::remove_file(staged).await;
422 Ok(())
423 }
424
425 async fn link(&self, content: &Path, final_path: &Path) -> Result<(), std::io::Error> {
426 let from = content.to_path_buf();
427 let to = final_path.to_path_buf();
428 let linked = tokio::task::spawn_blocking(move || std::fs::hard_link(&from, &to))
429 .await
430 .map_err(std::io::Error::other)?;
431
432 match linked {
433 Ok(()) => Ok(()),
434 Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => Ok(()),
435 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Err(error),
436 Err(_) => fs::copy(content, final_path).await.map(|_| ()),
440 }
441 }
442}