mtorrent_core/data/
storage.rs

1use crate::data::Error;
2use crate::pwp;
3use mtorrent_utils::warn_stopwatch;
4use sha1_smol::Sha1;
5use std::collections::BTreeMap;
6use std::path::{Path, PathBuf};
7use std::{cmp, fs, io};
8use tokio::sync::{mpsc, oneshot};
9
10/// Create new storage handle-actor pair, for files specified by `length_path_it` in the directory
11/// `parent_dir`. If the files don't exist, new files with the specified size will be created.
12pub fn new_async_storage(
13    parent_dir: impl AsRef<Path>,
14    length_path_it: impl Iterator<Item = (usize, PathBuf)>,
15) -> Result<(StorageClient, StorageServer), Error> {
16    let storage = Storage::new(parent_dir, length_path_it)?;
17    let (client, server) = async_generic_storage(storage);
18    Ok((client, StorageServer(server)))
19}
20
21/// Actor that performs filesystem operations, like reading and writing
22/// chunks of data, as well as calculating their hash. All operations are done
23/// sequentially and in the same order as they were scheduled.
24pub struct StorageServer(GenericStorageServer<fs::File>);
25
26impl StorageServer {
27    /// Start serving commands received from [`StorageClient`]. Filesystem operations will be
28    /// performed synchronously in the current thread.
29    pub async fn run(self) {
30        self.0.run().await;
31    }
32}
33
34/// Handle for requesting filesystem operations.
35#[derive(Clone)]
36pub struct StorageClient {
37    channel: mpsc::UnboundedSender<Command>,
38}
39
40impl StorageClient {
41    /// Write bytes `data` at `global_offset` relative to the start of the first file managed by
42    /// this storage. Returns once the filesystem write operation is finished.
43    pub async fn write_block(&self, global_offset: usize, data: Vec<u8>) -> Result<(), Error> {
44        let (result_sender, result_receiver) = oneshot::channel::<WriteResult>();
45        self.channel.send(Command::WriteBlock {
46            global_offset,
47            data,
48            callback: Some(result_sender),
49        })?;
50        result_receiver.await?
51    }
52
53    /// Schedule write of `data` at `global_offset` relative to the start of the first file managed
54    /// by this storage. Returns immediately without waiting for the result of the filesystem
55    /// operation.
56    pub fn start_write_block(&self, global_offset: usize, data: Vec<u8>) -> Result<(), Error> {
57        self.channel.send(Command::WriteBlock {
58            global_offset,
59            data,
60            callback: None,
61        })?;
62        Ok(())
63    }
64
65    /// Read `length` bytes at `global_offset` relative to the start of the first file managed by
66    /// this storage.
67    pub async fn read_block(&self, global_offset: usize, length: usize) -> Result<Vec<u8>, Error> {
68        let (result_sender, result_receiver) = oneshot::channel::<ReadResult>();
69        self.channel.send(Command::ReadBlock {
70            global_offset,
71            length,
72            callback: result_sender,
73        })?;
74        result_receiver.await?
75    }
76
77    /// Read `length` bytes at `global_offset` relative to the start of the first file managed by
78    /// this storage, calculate their SHA-1 hash, and compare it to `expected_sha`. Return
79    /// whether the two hashes are identical.
80    pub async fn verify_block(
81        &self,
82        global_offset: usize,
83        length: usize,
84        expected_sha1: &[u8; 20],
85    ) -> Result<bool, Error> {
86        let _sw = warn_stopwatch!("Verification of {length} bytes");
87        let (result_sender, result_receiver) = oneshot::channel::<VerifyResult>();
88        self.channel.send(Command::VerifyBlock {
89            global_offset,
90            length,
91            expected_sha1: *expected_sha1,
92            callback: result_sender,
93        })?;
94        result_receiver.await?
95    }
96}
97
98#[cfg(feature = "mocks")]
99#[doc(hidden)]
100pub fn new_mock_storage(total_size: usize) -> StorageClient {
101    let (tx, mut rx) = mpsc::unbounded_channel::<Command>();
102    tokio::task::spawn(async move {
103        while let Some(cmd) = rx.recv().await {
104            match cmd {
105                Command::WriteBlock {
106                    global_offset,
107                    data,
108                    callback,
109                } => {
110                    if let Some(cb) = callback {
111                        cb.send(if global_offset + data.len() < total_size {
112                            Ok(())
113                        } else {
114                            Err(Error::InvalidLocation)
115                        })
116                        .unwrap();
117                    }
118                }
119                Command::ReadBlock {
120                    global_offset,
121                    length,
122                    callback,
123                } => {
124                    callback
125                        .send(if global_offset + length < total_size {
126                            Ok(vec![0; length])
127                        } else {
128                            Err(Error::InvalidLocation)
129                        })
130                        .unwrap();
131                }
132                Command::VerifyBlock {
133                    global_offset,
134                    length,
135                    expected_sha1: _,
136                    callback,
137                } => {
138                    callback
139                        .send(if global_offset + length < total_size {
140                            Ok(true)
141                        } else {
142                            Err(Error::InvalidLocation)
143                        })
144                        .unwrap();
145                }
146            }
147        }
148    });
149    StorageClient { channel: tx }
150}
151
152// ------------------------------------------------------------------------------------------------
153
154fn async_generic_storage<F: RandomAccessReadWrite>(
155    storage: GenericStorage<F>,
156) -> (StorageClient, GenericStorageServer<F>) {
157    let (tx, rx) = mpsc::unbounded_channel::<Command>();
158    (
159        StorageClient { channel: tx },
160        GenericStorageServer {
161            channel: rx,
162            storage,
163        },
164    )
165}
166
167type WriteResult = Result<(), Error>;
168type ReadResult = Result<Vec<u8>, Error>;
169type VerifyResult = Result<bool, Error>;
170
171#[allow(clippy::enum_variant_names)]
172#[derive(Debug)]
173enum Command {
174    WriteBlock {
175        global_offset: usize,
176        data: Vec<u8>,
177        callback: Option<oneshot::Sender<WriteResult>>,
178    },
179    ReadBlock {
180        global_offset: usize,
181        length: usize,
182        callback: oneshot::Sender<ReadResult>,
183    },
184    VerifyBlock {
185        global_offset: usize,
186        length: usize,
187        expected_sha1: [u8; 20],
188        callback: oneshot::Sender<VerifyResult>,
189    },
190}
191
192struct GenericStorageServer<F: RandomAccessReadWrite> {
193    channel: mpsc::UnboundedReceiver<Command>,
194    storage: GenericStorage<F>,
195}
196
197impl<F: RandomAccessReadWrite> GenericStorageServer<F> {
198    async fn run(mut self) {
199        while let Some(cmd) = self.channel.recv().await {
200            self.handle_cmd(cmd);
201        }
202    }
203
204    fn handle_cmd(&self, cmd: Command) {
205        match cmd {
206            Command::WriteBlock {
207                global_offset,
208                data,
209                callback,
210            } => {
211                let result = self.storage.write_block(global_offset, data);
212                if let Some(callback) = callback {
213                    let _ = callback.send(result);
214                } else if let Err(e) = result {
215                    log::error!("Failed to write block: {e}");
216                }
217            }
218            Command::ReadBlock {
219                global_offset,
220                length,
221                callback,
222            } => {
223                let result = self.storage.read_block(global_offset, length);
224                let _ = callback.send(result);
225            }
226            Command::VerifyBlock {
227                global_offset,
228                length,
229                expected_sha1,
230                callback,
231            } => {
232                let end = global_offset + length;
233                let mut buffer = [0u8; pwp::MAX_BLOCK_SIZE];
234                let mut sha1 = Sha1::new();
235                let result = (global_offset..end)
236                    .step_by(buffer.len())
237                    .try_for_each(|offset| {
238                        let bytes_to_read = cmp::min(buffer.len(), end - offset);
239                        let dest = &mut buffer[..bytes_to_read];
240                        self.storage.read_block_into(offset, dest).map(|_| sha1.update(dest))
241                    })
242                    .map(|_| {
243                        let computed_sha1: [u8; 20] = sha1.digest().bytes();
244                        computed_sha1 == expected_sha1
245                    });
246                let _ = callback.send(result);
247            }
248        }
249    }
250}
251
252// ------------------------------------------------------------------------------------------------
253
254pub(super) type Storage = GenericStorage<fs::File>;
255
256pub(super) struct GenericStorage<F: RandomAccessReadWrite> {
257    files: BTreeMap<usize, F>,
258}
259
260impl Storage {
261    pub(super) fn new<I: Iterator<Item = (usize, PathBuf)>, P: AsRef<Path>>(
262        parent_dir: P,
263        length_path_it: I,
264    ) -> Result<Self, Error> {
265        let open_file = |(length, path): (usize, PathBuf)| -> io::Result<(usize, fs::File)> {
266            let path = parent_dir.as_ref().join(path);
267            if let Some(prefix) = path.parent() {
268                fs::create_dir_all(prefix)?;
269            }
270            let file = fs::OpenOptions::new()
271                .write(true)
272                .read(true)
273                .create(true)
274                .truncate(false)
275                .open(path)?;
276            file.set_len(length as u64)?;
277            Ok((length, file))
278        };
279
280        Self::from_length_file_pairs(length_path_it.map(open_file))
281    }
282}
283
284impl<F: RandomAccessReadWrite> GenericStorage<F> {
285    fn from_length_file_pairs<I: Iterator<Item = io::Result<(usize, F)>>>(
286        length_file_it: I,
287    ) -> Result<Self, Error> {
288        let mut filemap = BTreeMap::new();
289        let mut offset = 0usize;
290
291        for result in length_file_it {
292            let (length, file) = result?;
293            filemap.insert(offset, file);
294            offset += length;
295        }
296        if let Some((_offset, file)) = filemap.last_key_value() {
297            let fd_clone = file.try_clone()?;
298            filemap.insert(offset, fd_clone);
299        }
300        Ok(Self { files: filemap })
301    }
302
303    pub(super) fn write_block(&self, global_offset: usize, block: Vec<u8>) -> Result<(), Error> {
304        self.write_block_from(global_offset, &block)?;
305        Ok(())
306    }
307
308    pub(super) fn read_block(&self, global_offset: usize, length: usize) -> Result<Vec<u8>, Error> {
309        let mut dest = vec![0u8; length];
310        self.read_block_into(global_offset, &mut dest)?;
311        Ok(dest)
312    }
313
314    fn find_file_and_offset(&self, global_offset: usize) -> Result<(usize, &F, usize), Error> {
315        let next_start_offset = {
316            let (offset, _) =
317                self.files.range(global_offset + 1..).next().ok_or(Error::InvalidLocation)?;
318            *offset
319        };
320        let (start_offset, file) = {
321            let (offset, file) =
322                self.files.range(..=global_offset).last().ok_or(Error::InvalidLocation)?;
323            (*offset, file)
324        };
325        Ok((start_offset, file, next_start_offset))
326    }
327
328    fn write_block_from(&self, global_offset: usize, src: &[u8]) -> Result<(), Error> {
329        let (start_offset, file, next_start_offset) = self.find_file_and_offset(global_offset)?;
330        let local_offset = global_offset - start_offset;
331
332        let available_space = next_start_offset - global_offset;
333        if src.len() <= available_space {
334            file.write_all_at_offset(src, local_offset as u64)?;
335            Ok(())
336        } else {
337            let (left, right) = src.split_at(available_space);
338            file.write_all_at_offset(left, local_offset as u64)?;
339            self.write_block_from(next_start_offset, right)
340        }
341    }
342
343    fn read_block_into(&self, global_offset: usize, dest: &mut [u8]) -> Result<(), Error> {
344        let (start_offset, file, next_start_offset) = self.find_file_and_offset(global_offset)?;
345        let local_offset = global_offset - start_offset;
346
347        let available_space = next_start_offset - global_offset;
348        if dest.len() <= available_space {
349            file.read_all_at_offset(dest, local_offset as u64)?;
350            Ok(())
351        } else {
352            let (left, right) = dest.split_at_mut(available_space);
353            file.read_all_at_offset(left, local_offset as u64)?;
354            self.read_block_into(next_start_offset, right)
355        }
356    }
357}
358
359pub(super) trait RandomAccessReadWrite {
360    fn read_at_offset(&self, dest: &mut [u8], offset: u64) -> io::Result<usize>;
361    fn write_at_offset(&self, src: &[u8], offset: u64) -> io::Result<usize>;
362    fn try_clone(&self) -> io::Result<Self>
363    where
364        Self: Sized;
365
366    fn read_all_at_offset(&self, mut dest: &mut [u8], mut offset: u64) -> io::Result<()> {
367        while !dest.is_empty() {
368            let bytes_read = self.read_at_offset(dest, offset)?;
369            if bytes_read == 0 {
370                return Err(io::Error::new(
371                    io::ErrorKind::UnexpectedEof,
372                    "failed to fill whole buffer",
373                ));
374            }
375            dest = &mut dest[bytes_read..];
376            offset += bytes_read as u64;
377        }
378        Ok(())
379    }
380    fn write_all_at_offset(&self, mut src: &[u8], mut offset: u64) -> io::Result<()> {
381        while !src.is_empty() {
382            let bytes_written = self.write_at_offset(src, offset)?;
383            if bytes_written == 0 {
384                return Err(io::Error::new(
385                    io::ErrorKind::WriteZero,
386                    "failed to write whole buffer",
387                ));
388            }
389            src = &src[bytes_written..];
390            offset += bytes_written as u64;
391        }
392        Ok(())
393    }
394}
395
396#[cfg(unix)]
397impl RandomAccessReadWrite for fs::File {
398    fn read_at_offset(&self, dest: &mut [u8], offset: u64) -> io::Result<usize> {
399        use std::os::unix::prelude::*;
400        self.read_at(dest, offset)
401    }
402
403    fn write_at_offset(&self, src: &[u8], offset: u64) -> io::Result<usize> {
404        use std::os::unix::prelude::*;
405        self.write_at(src, offset)
406    }
407
408    fn try_clone(&self) -> io::Result<Self>
409    where
410        Self: Sized,
411    {
412        self.try_clone()
413    }
414}
415
416#[cfg(windows)]
417impl RandomAccessReadWrite for fs::File {
418    fn read_at_offset(&self, dest: &mut [u8], offset: u64) -> io::Result<usize> {
419        use std::os::windows::prelude::*;
420        self.seek_read(dest, offset)
421    }
422
423    fn write_at_offset(&self, src: &[u8], offset: u64) -> io::Result<usize> {
424        use std::os::windows::prelude::*;
425        self.seek_write(src, offset)
426    }
427
428    fn try_clone(&self) -> io::Result<Self>
429    where
430        Self: Sized,
431    {
432        self.try_clone()
433    }
434}
435
436#[cfg(test)]
437mod tests {
438    use super::*;
439    use std::io::{Cursor, Read, Seek, SeekFrom, Write};
440    use std::iter;
441    use tokio::task;
442
443    type FakeFile = std::cell::RefCell<Cursor<Vec<u8>>>;
444
445    fn fake_length_file_pair(content: Vec<u8>) -> io::Result<(usize, FakeFile)> {
446        Ok((content.len(), std::cell::RefCell::new(Cursor::new(content))))
447    }
448
449    impl RandomAccessReadWrite for FakeFile {
450        fn read_at_offset(&self, dest: &mut [u8], offset: u64) -> io::Result<usize> {
451            self.borrow_mut().seek(SeekFrom::Start(offset))?;
452            self.borrow_mut().read(dest)
453        }
454
455        fn write_at_offset(&self, src: &[u8], offset: u64) -> io::Result<usize> {
456            self.borrow_mut().seek(SeekFrom::Start(offset))?;
457            self.borrow_mut().write(src)
458        }
459
460        fn try_clone(&self) -> io::Result<Self>
461        where
462            Self: Sized,
463        {
464            Ok(self.clone())
465        }
466    }
467
468    #[test]
469    fn test_write_piece_within_one_file() {
470        let s = GenericStorage::from_length_file_pairs(
471            iter::repeat_with(|| fake_length_file_pair(vec![0u8; 10])).take(3),
472        )
473        .unwrap();
474
475        s.write_block(16, vec![1u8, 2u8, 3u8, 4u8]).unwrap();
476
477        assert_eq!(&vec![0u8; 10], s.files.get(&0).unwrap().borrow().get_ref());
478        assert_eq!(
479            &vec![0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 1u8, 2u8, 3u8, 4u8],
480            s.files.get(&10).unwrap().borrow().get_ref()
481        );
482        assert_eq!(&vec![0u8; 10], s.files.get(&20).unwrap().borrow().get_ref());
483    }
484
485    #[test]
486    fn test_write_piece_on_file_boundary() {
487        let s = GenericStorage::from_length_file_pairs(
488            iter::repeat_with(|| fake_length_file_pair(vec![0u8; 10])).take(3),
489        )
490        .unwrap();
491
492        s.write_block(17, vec![1u8, 2u8, 3u8, 4u8, 5u8]).unwrap();
493
494        assert_eq!(&vec![0u8; 10], s.files.get(&0).unwrap().borrow().get_ref());
495        assert_eq!(
496            &vec![0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 1u8, 2u8, 3u8],
497            s.files.get(&10).unwrap().borrow().get_ref()
498        );
499        assert_eq!(
500            &vec![4u8, 5u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8],
501            s.files.get(&20).unwrap().borrow().get_ref()
502        );
503    }
504
505    #[test]
506    fn test_read_piece_within_one_file() {
507        let s = GenericStorage::from_length_file_pairs(
508            iter::repeat_with(|| fake_length_file_pair((0u8..10u8).collect())).take(3),
509        )
510        .unwrap();
511
512        let dest = s.read_block(12, 3).unwrap();
513        assert_eq!(vec![2u8, 3u8, 4u8], dest);
514
515        assert_eq!(&(0u8..10u8).collect::<Vec<u8>>(), s.files.get(&0).unwrap().borrow().get_ref());
516        assert_eq!(&(0u8..10u8).collect::<Vec<u8>>(), s.files.get(&10).unwrap().borrow().get_ref());
517        assert_eq!(&(0u8..10u8).collect::<Vec<u8>>(), s.files.get(&20).unwrap().borrow().get_ref());
518    }
519
520    #[test]
521    fn test_read_piece_on_file_boundary() {
522        let s = GenericStorage::from_length_file_pairs(
523            iter::repeat_with(|| fake_length_file_pair((1u8..=10u8).collect())).take(2),
524        )
525        .unwrap();
526
527        let dest = s.read_block(8, 3).unwrap();
528        assert_eq!(vec![9u8, 10u8, 1u8], dest);
529
530        assert_eq!(&(1u8..=10u8).collect::<Vec<u8>>(), s.files.get(&0).unwrap().borrow().get_ref());
531        assert_eq!(
532            &(1u8..=10u8).collect::<Vec<u8>>(),
533            s.files.get(&10).unwrap().borrow().get_ref()
534        );
535    }
536
537    #[test]
538    fn test_past_the_end_read_fails() {
539        let s = GenericStorage::from_length_file_pairs(iter::once(fake_length_file_pair(
540            (1u8..=10u8).collect(),
541        )))
542        .unwrap();
543        assert!(matches!(s.read_block(5, 10), Err(Error::InvalidLocation)));
544        assert!(matches!(s.read_block(11, 5), Err(Error::InvalidLocation)));
545    }
546
547    #[tokio::test]
548    async fn test_async_detached_write_then_read_on_file_boundary() {
549        task::LocalSet::new()
550            .run_until(async {
551                let s = GenericStorage::from_length_file_pairs(
552                    iter::repeat_with(|| fake_length_file_pair(vec![0u8; 10])).take(2),
553                )
554                .unwrap();
555
556                let (client, server) = async_generic_storage(s);
557
558                task::spawn_local(async move {
559                    server.run().await;
560                });
561
562                // given
563                let initial_data = client.read_block(8, 3).await.unwrap();
564                assert_eq!(vec![0u8, 0u8, 0u8], initial_data);
565
566                // when
567                client.start_write_block(8, vec![9u8, 10u8, 1u8]).unwrap();
568
569                // then
570                let final_data = client.read_block(8, 3).await.unwrap();
571                assert_eq!(vec![9u8, 10u8, 1u8], final_data);
572            })
573            .await;
574    }
575
576    #[tokio::test]
577    async fn test_async_verify_block() {
578        let sha1_0_10 =
579            b"\x49\x41\x79\x71\x4a\x6c\xd6\x27\x23\x9d\xfe\xde\xdf\x2d\xe9\xef\x99\x4c\xaf\x03";
580        let sha1_10_20 =
581            b"\xdd\xd1\x27\x8d\x28\xaf\x87\xc7\x58\x84\xf5\x5b\x71\xfb\xb4\xa1\x23\x1a\xf2\xe5";
582        task::LocalSet::new()
583            .run_until(async {
584                let s = GenericStorage::from_length_file_pairs(iter::once(fake_length_file_pair(
585                    (0u8..20u8).collect(),
586                )))
587                .unwrap();
588
589                let (client, server) = async_generic_storage(s);
590
591                task::spawn_local(async move {
592                    server.run().await;
593                });
594
595                let verify_success = client.verify_block(0, 10, sha1_0_10).await.unwrap();
596                assert!(verify_success);
597
598                let verify_success = client.verify_block(0, 10, sha1_10_20).await.unwrap();
599                assert!(!verify_success);
600
601                let verify_success = client.verify_block(10, 10, sha1_10_20).await.unwrap();
602                assert!(verify_success);
603
604                let verify_success = client.verify_block(10, 10, sha1_0_10).await.unwrap();
605                assert!(!verify_success);
606            })
607            .await;
608    }
609}
mtorrent_core/data/storage.rs

mtorrent_core/data/
storage.rs
