Skip to main content

lfsx_server/storage/
mod.rs

1mod sweep;
2
3use std::path::{Path, PathBuf};
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::time::Instant;
6
7use tokio::sync::Mutex;
8
9use futures_util::{Stream, StreamExt};
10use sha2::{Digest, Sha256};
11use tokio::fs;
12use tokio::io::AsyncWriteExt;
13
14use crate::error::Error;
15use crate::namespace::Namespace;
16
17pub use sweep::SweepReport;
18
19pub struct LocalStore {
20    root: PathBuf,
21    counter: AtomicU64,
22    usage: Mutex<Option<(Instant, u64, u64)>>,
23    scans: AtomicU64,
24}
25
26impl LocalStore {
27    pub fn new(root: impl Into<PathBuf>) -> Self {
28        Self {
29            root: root.into(),
30            counter: AtomicU64::new(0),
31            usage: Mutex::new(None),
32            scans: AtomicU64::new(0),
33        }
34    }
35
36    pub fn validate_oid(oid: &str) -> Result<(), Error> {
37        let well_formed = oid.len() == 64
38            && oid
39                .bytes()
40                .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b));
41
42        well_formed.then_some(()).ok_or(Error::MalformedOid)
43    }
44
45    fn object_path(&self, ns: &Namespace, oid: &str) -> PathBuf {
46        self.root
47            .join(ns.org())
48            .join(ns.repo())
49            .join(&oid[0..2])
50            .join(&oid[2..4])
51            .join(oid)
52    }
53
54    pub fn scans(&self) -> u64 {
55        self.scans.load(Ordering::Relaxed)
56    }
57
58    pub async fn writable(&self) -> Result<(), Error> {
59        fs::create_dir_all(&self.root).await?;
60
61        let ticket = self.counter.fetch_add(1, Ordering::Relaxed);
62        let probe = self.root.join(format!(".readiness.{ticket}"));
63
64        fs::write(&probe, b"").await?;
65        fs::remove_file(&probe).await?;
66
67        Ok(())
68    }
69
70    pub async fn exists(&self, ns: &Namespace, oid: &str) -> bool {
71        Self::validate_oid(oid).is_ok() && fs::metadata(self.object_path(ns, oid)).await.is_ok()
72    }
73
74    pub async fn open(&self, ns: &Namespace, oid: &str) -> Result<(fs::File, u64), Error> {
75        Self::validate_oid(oid)?;
76        let path = self.object_path(ns, oid);
77        let file = fs::File::open(&path).await.map_err(|_| Error::NotFound)?;
78        let size = file.metadata().await?.len();
79        Ok((file, size))
80    }
81
82    pub async fn write<S, E>(
83        &self,
84        ns: &Namespace,
85        oid: &str,
86        expected_size: Option<u64>,
87        mut chunks: S,
88    ) -> Result<u64, Error>
89    where
90        S: Stream<Item = Result<axum::body::Bytes, E>> + Unpin,
91        E: std::error::Error + Send + Sync + 'static,
92    {
93        Self::validate_oid(oid)?;
94
95        let path = self.object_path(ns, oid);
96        let parent = path.parent().expect("object paths always have a parent");
97        fs::create_dir_all(parent).await?;
98
99        let staged = self.staging_path(parent, oid);
100        let outcome = self.stream_to(&staged, &mut chunks).await;
101
102        match outcome {
103            Ok((digest, written)) => {
104                self.finish(&staged, &path, oid, expected_size, &digest, written)
105                    .await?;
106                Ok(written)
107            }
108            Err(error) => {
109                let _ = fs::remove_file(&staged).await;
110                Err(error)
111            }
112        }
113    }
114
115    fn staging_path(&self, parent: &Path, oid: &str) -> PathBuf {
116        let ticket = self.counter.fetch_add(1, Ordering::Relaxed);
117        parent.join(format!("{oid}.{ticket}.part"))
118    }
119
120    async fn stream_to<S, E>(&self, staged: &Path, chunks: &mut S) -> Result<(String, u64), Error>
121    where
122        S: Stream<Item = Result<axum::body::Bytes, E>> + Unpin,
123        E: std::error::Error + Send + Sync + 'static,
124    {
125        let mut file = fs::File::create(staged).await?;
126        let mut hasher = Sha256::new();
127        let mut written = 0u64;
128
129        while let Some(chunk) = chunks.next().await {
130            let chunk = chunk.map_err(std::io::Error::other)?;
131            hasher.update(&chunk);
132            written += chunk.len() as u64;
133            file.write_all(&chunk).await?;
134        }
135
136        file.flush().await?;
137        file.sync_all().await?;
138
139        Ok((hex::encode(hasher.finalize()), written))
140    }
141
142    async fn finish(
143        &self,
144        staged: &Path,
145        final_path: &Path,
146        oid: &str,
147        expected_size: Option<u64>,
148        digest: &str,
149        written: u64,
150    ) -> Result<(), Error> {
151        if let Some(declared) = expected_size.filter(|declared| *declared != written) {
152            let _ = fs::remove_file(staged).await;
153            return Err(Error::SizeMismatch {
154                declared,
155                actual: written,
156            });
157        }
158
159        if digest != oid {
160            let _ = fs::remove_file(staged).await;
161            return Err(Error::OidMismatch {
162                declared: oid.to_owned(),
163                actual: digest.to_owned(),
164            });
165        }
166
167        fs::rename(staged, final_path).await?;
168        Ok(())
169    }
170}