Skip to main content

lfsx_server/storage/
mod.rs

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
59// A transfer that has passed every check and is waiting to be put somewhere.
60pub struct Staged {
61    pub path: PathBuf,
62    destination: PathBuf,
63    pub written: u64,
64    fresh: bool,
65}
66
67// What a download reads from, whether or not the bytes on disk are the object.
68pub 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// What is left of a repository's budget for one transfer. It travels with the
124// upload because a client that skips negotiation may also skip declaring a
125// size, and a budget checked once against a number the client chose is not a
126// budget.
127#[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    // Everything a transfer has to survive before it counts as an object: the
241    // digest it claims, the size it declared, the ceiling on a single object and
242    // the repository's remaining budget. It ends on local disk whatever the
243    // backend is, because a bucket cannot be asked to hold bytes that might turn
244    // out to be the wrong ones.
245    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        // A retried transfer of an object this repository already holds costs it
270        // no room, so it must not count against the budget a second time.
271        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        // The digest, the declared size and the budget are all counted on the
359        // plaintext going past, whatever the bytes look like once they land —
360        // so compression is a different sink, not a different path.
361        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            // The declared size is a claim by the client, so the ceiling has to
374            // hold against a body that ignores it. Stopping at the chunk that
375            // crosses the line is the point: reading to the end to find out how
376            // big it was would be the outage this limit exists to prevent.
377            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    // One copy of the bytes under .content, and a hard link per repository that
394    // holds them. Two projects sharing an asset pack cost the disk once, and the
395    // link count is the reference count — the filesystem does the bookkeeping, so
396    // nothing can leak a repository's contents to another and nothing needs a
397    // migration: objects already sitting at their repository path keep working as
398    // ordinary files with one link.
399    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            // The content was collected between finding it and linking to it:
410            // a concurrent retain on another repository dropped its last other
411            // reference. The staged copy is still here precisely for this, so
412            // put it back and link again rather than failing a push that did
413            // nothing wrong.
414            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            // A filesystem without hard links, or one crossing a device
437            // boundary: fall back to a full copy so the transfer still
438            // succeeds. The disk pays for it, the client never notices.
439            Err(_) => fs::copy(content, final_path).await.map(|_| ()),
440        }
441    }
442}