Skip to main content

webdav_server/storage/
transaction.rs

1use crate::storage::{
2    Error::{CreateTempFile, InvalidPath, RenameError, WriteChunkFailure},
3    Payload, Result,
4    file::Metadata,
5};
6use blake3::Hasher;
7use bytes::Bytes;
8use fs4::AsyncFileExt;
9use futures::{Stream, StreamExt};
10use std::{
11    error::Error as StdErr,
12    io::{Error as IoErr, ErrorKind::UnexpectedEof},
13    path::{Path, PathBuf},
14};
15use tokio::{
16    fs::{File, remove_file, rename},
17    io::AsyncWriteExt,
18};
19use uuid::Uuid;
20
21#[derive(Debug)]
22/// A single write transaction for committing a blob to the vault.
23///
24/// A `Transaction` is created by [`Service::try_save`] and represents the lifecycle
25/// of one upload from start to finish. It manages two files:
26///
27/// 1. A temporary staging file at `<vault>/<uuid>.tmp`, where bytes are streamed.
28/// 2. The final blob file at `<vault>/<blake3_hash>`, which is created via an atomic `rename(2)`.
29///
30/// If the transaction fails at any point, the `.tmp` file is deleted before the error
31/// is returned. This ensures the vault always contains only successfully committed blobs.
32pub(crate) struct Transaction {
33    temp: PathBuf,
34    _uuid: Uuid,
35    hasher: Hasher,
36}
37
38/// Constructors and accessors.
39#[allow(unused)]
40impl Transaction {
41    /// Creates a new transaction rooted in the given vault directory.
42    ///
43    /// A new UUID is generated on each call to ensure the staging file path
44    /// is unique, which prevents collisions between concurrent uploads.
45    pub(crate) fn new<P>(vault_path: P) -> Self
46    where
47        PathBuf: From<P>,
48    {
49        let uuid = Uuid::new_v4();
50        let vault_path: PathBuf = vault_path.into();
51        let mut temp = vault_path.join(uuid.to_string());
52        temp.add_extension("tmp");
53        Self {
54            temp,
55            _uuid: uuid,
56            hasher: Hasher::new(),
57        }
58    }
59
60    /// Returns the path of the temporary staging file.
61    ///
62    /// Primarily exposed for tests that want to assert cleanup happened.
63    pub(crate) fn temp_path(&self) -> &Path {
64        &self.temp
65    }
66}
67
68impl Transaction {
69    /// Streams the payload to disk and commits it atomically.
70    ///
71    /// This is the public entry point. It delegates to `write_and_commit` and
72    /// guarantees that the temporary staging file is deleted if anything goes wrong.
73    ///
74    /// # Errors
75    /// See [`Transaction::write_and_commit`] for the full list of failure modes.
76    pub async fn commit<S, E>(
77        mut self,
78        payload: Payload<S, E>,
79    ) -> Result<Metadata>
80    where
81        S: Stream<Item = std::result::Result<Bytes, E>> + Unpin,
82        E: Into<Box<dyn StdErr + Send + Sync>>,
83    {
84        match self.write_and_commit(payload).await {
85            Ok(metadata) => Ok(metadata),
86
87            Err(e) => {
88                // Clean up the staging file before propagating the error.
89                // If the cleanup itself fails, we silently ignore it so the
90                // original error `e` is not lost.
91                //
92                // Future: instead of deleting, queue for GC so interrupted
93                // uploads can be resumed (tus-style resumable uploads).
94                let _ = remove_file(self.temp).await;
95
96                Err(e)
97            }
98        }
99    }
100
101    /// The internal implementation of the write pipeline.
102    ///
103    /// Steps:
104    /// 1. Create the `.tmp` staging file.
105    /// 2. Pre-allocate the expected number of bytes on disk via `fallocate`
106    ///    (skipped for zero-byte streams). This reduces fragmentation on
107    ///    spinning drives and SD cards.
108    /// 3. Stream all chunks to disk, hashing each one incrementally with BLAKE3.
109    /// 4. Validate that the bytes written match `expected_size` (EOF check).
110    /// 5. Atomically rename the `.tmp` file to `<vault>/<blake3_hash>`.
111    ///
112    /// # Errors
113    /// - [`Error::CreateTempFile`] if the staging file cannot be created.
114    /// - [`Error::StreamReadError`] if a chunk cannot be read from the network stream.
115    /// - [`Error::WriteChunkFailure`] if a chunk cannot be written to disk.
116    /// - [`Error::StreamReadError`] if `bytes_written != expected_size` at EOF.
117    /// - [`Error::InvalidPath`] if the staging file has no parent directory.
118    /// - [`Error::RenameError`] if the atomic rename fails.
119    ///
120    /// [`Error::CreateTempFile`]: crate::storage::Error::CreateTempFile
121    /// [`Error::StreamReadError`]: crate::storage::Error::StreamReadError
122    /// [`Error::WriteChunkFailure`]: crate::storage::Error::WriteChunkFailure
123    /// [`Error::InvalidPath`]: crate::storage::Error::InvalidPath
124    /// [`Error::RenameError`]: crate::storage::Error::RenameError
125    async fn write_and_commit<S, E>(
126        &mut self,
127        payload: Payload<S, E>,
128    ) -> Result<Metadata>
129    where
130        S: Stream<Item = std::result::Result<Bytes, E>> + Unpin,
131        E: Into<Box<dyn StdErr + Send + Sync>>,
132    {
133        let mut file =
134            File::create(&self.temp).await.map_err(|e| CreateTempFile {
135                path: self.temp.clone(),
136                source: e,
137            })?;
138
139        if payload.expected_size > 0 {
140            let _ = file.allocate(payload.expected_size).await;
141        }
142
143        let bytes_written =
144            self.process_stream(payload.stream, &mut file).await?;
145
146        if bytes_written != payload.expected_size {
147            return Err(IoErr::new(
148                UnexpectedEof,
149                "Content-Length doesn't match the bytes streamed",
150            )
151            .into());
152        }
153
154        let target = self
155            .temp
156            .parent()
157            .ok_or_else(|| InvalidPath {
158                path: self.temp.clone(),
159            })?
160            .join(self.hasher.finalize().to_string());
161
162        rename(&self.temp, &target).await.map_err(|e| RenameError {
163            path: target,
164            source: e,
165        })?;
166
167        Metadata::try_new(&file.metadata().await?, self.hasher.finalize())
168    }
169
170    /// Reads all chunks from the stream, writes them to disk, and returns the total bytes written.
171    ///
172    /// Each chunk is also fed to the BLAKE3 hasher incrementally, so there is no
173    /// need to re-read the file after streaming to compute the hash.
174    ///
175    /// # Errors
176    /// - [`Error::StreamReadError`] if the network stream yields an error on a chunk.
177    /// - [`Error::WriteChunkFailure`] if writing a chunk to disk fails.
178    ///
179    /// [`Error::StreamReadError`]: crate::storage::Error::StreamReadError
180    /// [`Error::WriteChunkFailure`]: crate::storage::Error::WriteChunkFailure
181    async fn process_stream<S, E>(
182        &mut self,
183        mut f_stream: S,
184        file: &mut File,
185    ) -> Result<u64>
186    where
187        S: Stream<Item = std::result::Result<Bytes, E>> + Unpin,
188        E: Into<Box<dyn StdErr + Send + Sync>>,
189    {
190        let mut bytes_written: u64 = 0;
191
192        // The addition here is safe: the upload handler rejects payloads larger
193        // than 10 GB, so `bytes_written` can never overflow a u64.
194        #[allow(clippy::arithmetic_side_effects)]
195        while let Some(chunk) = f_stream.next().await {
196            let chunk = chunk.map_err(|e| IoErr::other(e))?;
197
198            file.write_all(&chunk)
199                .await
200                .map_err(|e| WriteChunkFailure {
201                    path: self.temp.clone(),
202                    source: e,
203                })?;
204
205            bytes_written += u64::try_from(chunk.len())?;
206
207            self.hasher.update(&chunk);
208        }
209        Ok(bytes_written)
210    }
211}
212
213impl AsRef<Self> for Transaction {
214    fn as_ref(&self) -> &Self {
215        self
216    }
217}
218
219#[cfg(test)]
220mod test {
221    use blake3::Hasher;
222    use bytes::Bytes;
223    use std::{
224        io::{Error as IoErr, ErrorKind},
225        path::PathBuf,
226    };
227
228    use crate::storage::{
229        Payload, tests::with_temp_transaction, transaction::Transaction,
230    };
231
232    #[tokio::test]
233    async fn successful_commit_and_hash() {
234        with_temp_transaction(async move |transaction, vault_path| {
235            let chunks: Vec<Result<Bytes, IoErr>> = vec![
236                Ok(Bytes::from("hello")),
237                Ok(Bytes::from(" ")),
238                Ok(Bytes::from("world")),
239            ];
240
241            let payload =
242                Payload::new(11 as u64, futures::stream::iter(chunks));
243
244            let result = transaction.commit(payload).await;
245
246            assert!(result.is_ok());
247
248            let metadata = result.as_ref().unwrap();
249
250            let target_path = vault_path.join(metadata.hash.to_string());
251
252            let mut hasher = Hasher::new();
253            let bytes = tokio::fs::read(target_path).await.unwrap();
254            hasher.update(&bytes);
255
256            let expected_hash = hasher.finalize().to_string();
257
258            assert_eq!(expected_hash, metadata.hash.to_string());
259        })
260        .await
261    }
262
263    #[tokio::test]
264    async fn zero_byte_stream_creates_empty_file() {
265        with_temp_transaction(async move |transaction, vault_path| {
266            let chunks: Vec<Result<Bytes, IoErr>> = Vec::new();
267
268            let payload = Payload::new(0 as u64, futures::stream::iter(chunks));
269
270            let result = transaction.commit(payload).await;
271
272            // Test: Should succeed w/o panic
273            assert!(result.is_ok());
274
275            let metadata = result.unwrap();
276
277            assert_eq!(metadata.size, 0);
278
279            let target_path = vault_path.join(metadata.hash.to_string());
280
281            // Test: There File should be present, even though its empty
282            assert!(target_path.exists());
283
284            let expected_hash = Hasher::new().finalize().to_string();
285
286            // Test: Hashes match
287            assert_eq!(metadata.hash.to_string(), expected_hash)
288        })
289        .await;
290    }
291
292    #[tokio::test]
293    async fn aborted_test_cleans_up_garbage() {
294        with_temp_transaction(async move |transaction, _vault_path| {
295            let temp_path = transaction.temp_path().to_owned();
296
297            let chunks: Vec<Result<Bytes, IoErr>> = vec![
298                Ok(Bytes::from("good bytes")),
299                Err(IoErr::new(ErrorKind::ConnectionAborted, "Wifi dies, lol")),
300            ];
301
302            let payload =
303                Payload::new(20 as u64, futures::stream::iter(chunks));
304
305            let result = transaction.commit(payload).await;
306
307            assert!(result.is_err());
308
309            assert!(!temp_path.exists());
310            // assert!(!target_path.exists());
311        })
312        .await;
313    }
314
315    #[tokio::test]
316    async fn transaction_fails_if_vault_missing() {
317        let vault = PathBuf::from("/tmp/path/that/possibly/doesnt/exist/lol");
318        let transaction = Transaction::new(vault);
319
320        let chunks: Vec<Result<Bytes, IoErr>> =
321            vec![Ok(Bytes::from("data_data"))];
322        let f_stream = futures::stream::iter(chunks);
323
324        let payload = Payload::new(9u64, f_stream);
325
326        let result = transaction.commit(payload).await;
327
328        // Test: No problem parsing the data
329        assert!(result.is_err());
330
331        use crate::storage::Error::CreateTempFile;
332
333        // Test: Yields CreateTempFile Error, 'cause vault directory was missing
334        assert!(
335            matches!(result, Err(CreateTempFile { .. })),
336            "Expected Err(CreateTempFile)"
337        );
338    }
339
340    #[tokio::test]
341    async fn hardcoded_hash_correctness() {
342        with_temp_transaction(async move |transaction, _vault_path| {
343            let payload : Vec<Result<Bytes, IoErr>> = vec![Ok(Bytes::from("hello world"))];
344            let f_stream = futures::stream::iter(payload);
345
346            let payload = Payload::new(11 as u64, f_stream);
347
348            let metadata = transaction.commit(payload).await.unwrap();
349
350            // Pre-calculated Blake3 hash of "hello world"
351            let expected_hash = "d74981efa70a0c880b8d8c1985d075dbcbf679b99a5f9914e5aaf96b831a9e24";
352
353            // Test: committed payload generates same hash as expected_hash
354            assert_eq!(metadata.hash.to_string(), expected_hash);
355        }).await;
356    }
357
358    #[tokio::test]
359    async fn mismatch_content_fails_plus_cleans_up() {
360        with_temp_transaction(async move |transaction, _| {
361            let temp_path = transaction.temp_path().to_owned();
362
363            let chunks: Vec<Result<Bytes, IoErr>> =
364                vec![Ok(Bytes::from("Halo there"))];
365
366            let payload = Payload::new(67u64, futures::stream::iter(chunks));
367
368            let result = transaction.commit(payload).await;
369
370            // Test: Unexpected EOF causes failure
371            assert!(result.is_err());
372
373            // Test: Cleanup is expected on failure
374            assert!(!temp_path.exists());
375        })
376        .await;
377    }
378}