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)]
22pub(crate) struct Transaction {
33 temp: PathBuf,
34 _uuid: Uuid,
35 hasher: Hasher,
36}
37
38#[allow(unused)]
40impl Transaction {
41 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 pub(crate) fn temp_path(&self) -> &Path {
64 &self.temp
65 }
66}
67
68impl Transaction {
69 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 let _ = remove_file(self.temp).await;
95
96 Err(e)
97 }
98 }
99 }
100
101 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 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 #[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 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 assert!(target_path.exists());
283
284 let expected_hash = Hasher::new().finalize().to_string();
285
286 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 })
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 assert!(result.is_err());
330
331 use crate::storage::Error::CreateTempFile;
332
333 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 let expected_hash = "d74981efa70a0c880b8d8c1985d075dbcbf679b99a5f9914e5aaf96b831a9e24";
352
353 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 assert!(result.is_err());
372
373 assert!(!temp_path.exists());
375 })
376 .await;
377 }
378}