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