virtual_file_transfer/
source_reader.rs

1use std::io::{self, Read, Seek, SeekFrom};
2
3use crate::source_reader::ReaderKind::{Seekable, Streaming};
4
5/// Supplies bytes to the receiver as either a sequential or seekable stream.
6///
7/// Use [`Self::streaming`] for `Read`-only sources, [`Self::seekable`] for
8/// `Read + Seek` sources, or `Self::spooled` (feature `spool`) to materialize
9/// a source into temporary storage.
10///
11/// The selected reader kind determines whether `IStream::Seek` can succeed
12/// and whether the stream size can be derived when not explicitly provided.
13pub struct SourceReader {
14    inner: ReaderKind,
15}
16
17/// Internal trait alias-like bound for readers that support both `Read`
18/// and `Seek`.
19pub trait ReadSeek: Read + Seek {}
20
21impl<T> ReadSeek for T where T: Read + Seek {}
22
23enum ReaderKind {
24    Streaming(Box<dyn Read + Send>),
25    Seekable(Box<dyn ReadSeek + Send>),
26}
27
28#[derive(Debug)]
29pub(crate) enum SourceReaderError {
30    NotSeekable,
31    Io,
32}
33
34impl SourceReader {
35    /// Wraps a sequential `Read` source without materializing its contents.
36    ///
37    /// The reader does not support seeking, and its length cannot be derived.
38    /// Supply a size using [`crate::VirtualFile::with_size`]
39    /// when the receiving application requires one.
40    pub fn streaming<R>(reader: R) -> Self
41    where
42        R: Read + Send + 'static,
43    {
44        Self {
45            inner: Streaming(Box::new(reader)),
46        }
47    }
48
49    /// Wraps a `Read + Seek` source that supports random access.
50    ///
51    /// When a file size is not provided, the length can be derived by seeking to
52    /// the end and restoring the original position. The underlying reader must
53    /// support those operations reliably.
54    pub fn seekable<R>(reader: R) -> Self
55    where
56        R: Read + Seek + Send + 'static,
57    {
58        Self {
59            inner: Seekable(Box::new(reader)),
60        }
61    }
62
63    /// Copies the complete source into a temporary file and returns a seekable
64    /// reader backed by that file.
65    ///
66    /// Spooling is explicit and finishes before this method returns. It trades
67    /// temporary disk usage and startup latency for seekability and a derivable
68    /// size. The temporary file's lifetime is managed with the returned reader.
69    ///
70    /// Requires the `spool` feature.
71    ///
72    /// # Errors
73    ///
74    /// Returns an I/O error if creating, writing, seeking, or reading the
75    /// temporary file or source fails.
76    #[cfg(feature = "spool")]
77    pub fn spooled<R>(reader: R) -> io::Result<Self>
78    where
79        R: Read + Send + 'static,
80    {
81        Ok(Self::seekable(spool_to_file(reader)?))
82    }
83
84    pub(crate) fn seek(&mut self, position: SeekFrom) -> Result<u64, SourceReaderError> {
85        match &mut self.inner {
86            ReaderKind::Seekable(reader) => {
87                reader.seek(position).map_err(|_| SourceReaderError::Io)
88            }
89
90            ReaderKind::Streaming(_) => Err(SourceReaderError::NotSeekable),
91        }
92    }
93
94    pub(crate) fn try_size(&mut self) -> io::Result<Option<u64>> {
95        match &mut self.inner {
96            ReaderKind::Streaming(_) => Ok(None),
97
98            ReaderKind::Seekable(reader) => {
99                let current = reader.stream_position()?;
100                let size = reader.seek(SeekFrom::End(0))?;
101                reader.seek(SeekFrom::Start(current))?;
102
103                Ok(Some(size))
104            }
105        }
106    }
107}
108
109impl Read for SourceReader {
110    fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
111        match &mut self.inner {
112            ReaderKind::Streaming(reader) => reader.read(buffer),
113            ReaderKind::Seekable(reader) => reader.read(buffer),
114        }
115    }
116}
117
118#[cfg(feature = "spool")]
119fn spool_to_file(mut reader: impl Read) -> io::Result<std::fs::File> {
120    let mut file = tempfile::tempfile()?;
121
122    io::copy(&mut reader, &mut file)?;
123    file.seek(SeekFrom::Start(0))?;
124
125    Ok(file)
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131
132    #[test]
133    fn streaming_reader_is_not_seekable() {
134        let mut reader = SourceReader::streaming(std::io::Cursor::new(b"hello".to_vec()));
135
136        assert!(matches!(
137            reader.seek(SeekFrom::Start(0)),
138            Err(SourceReaderError::NotSeekable)
139        ));
140    }
141
142    #[test]
143    fn streaming_reader_cannot_derive_size() {
144        let mut reader = SourceReader::streaming(std::io::Cursor::new(b"hello".to_vec()));
145
146        assert_eq!(reader.try_size().unwrap(), None);
147    }
148
149    #[test]
150    fn seekable_reader_can_seek() {
151        let mut reader = SourceReader::seekable(std::io::Cursor::new(b"hello".to_vec()));
152
153        assert_eq!(reader.seek(SeekFrom::Start(2)).unwrap(), 2);
154
155        let mut buf = [0u8; 2];
156        reader.read_exact(&mut buf).unwrap();
157
158        assert_eq!(&buf, b"ll");
159    }
160
161    #[test]
162    fn seekable_reader_can_derive_size() {
163        let mut reader = SourceReader::seekable(std::io::Cursor::new(b"hello".to_vec()));
164
165        assert_eq!(reader.try_size().unwrap(), Some(5));
166    }
167
168    #[test]
169    fn try_size_restores_position() {
170        let mut reader = SourceReader::seekable(std::io::Cursor::new(b"hello".to_vec()));
171
172        reader.seek(SeekFrom::Start(2)).unwrap();
173
174        assert_eq!(reader.try_size().unwrap(), Some(5));
175        assert_eq!(reader.seek(SeekFrom::Current(0)).unwrap(), 2);
176    }
177
178    #[cfg(feature = "spool")]
179    #[test]
180    fn spooled_reader_preserves_data_and_supports_seek() {
181        let data: Vec<u8> = (0..256 * 1024).map(|i| (i % 251) as u8).collect();
182
183        let mut reader = SourceReader::spooled(io::Cursor::new(data.clone())).unwrap();
184
185        // Size can be derived.
186        assert_eq!(reader.try_size().unwrap(), Some(data.len() as u64));
187
188        // Reading starts at position zero.
189        let mut actual = Vec::new();
190        reader.read_to_end(&mut actual).unwrap();
191        assert_eq!(actual, data);
192
193        // Random access is supported.
194        assert_eq!(reader.seek(SeekFrom::Start(42)).unwrap(), 42);
195
196        // Size lookup preserves the current position.
197        assert_eq!(reader.try_size().unwrap(), Some(data.len() as u64));
198        assert_eq!(reader.seek(SeekFrom::Current(0)).unwrap(), 42);
199
200        let mut buf = [0u8; 4];
201        reader.read_exact(&mut buf).unwrap();
202        assert_eq!(&buf, &data[42..46]);
203    }
204
205    #[cfg(feature = "spool")]
206    #[test]
207    fn spooled_reader_accepts_empty_source() {
208        let mut reader = SourceReader::spooled(io::empty()).unwrap();
209
210        assert_eq!(reader.try_size().unwrap(), Some(0));
211
212        let mut buf = [0u8; 1];
213        assert_eq!(reader.read(&mut buf).unwrap(), 0);
214    }
215
216    #[cfg(feature = "spool")]
217    #[test]
218    fn spooled_reader_propagates_source_error() {
219        struct FailingReader;
220
221        impl Read for FailingReader {
222            fn read(&mut self, _: &mut [u8]) -> io::Result<usize> {
223                Err(io::Error::other("source read failed"))
224            }
225        }
226
227        let result = SourceReader::spooled(FailingReader);
228
229        assert!(matches!(
230            result,
231            Err(ref err) if err.kind() == io::ErrorKind::Other
232        ));
233    }
234
235    #[cfg(feature = "spool")]
236    #[test]
237    fn spool_creates_file_with_correct_contents() {
238        let data = b"hello spool";
239
240        let mut file = spool_to_file(io::Cursor::new(data)).unwrap();
241
242        assert_eq!(file.metadata().unwrap().len(), data.len() as u64);
243
244        let mut actual = Vec::new();
245        file.read_to_end(&mut actual).unwrap();
246
247        assert_eq!(actual, data);
248
249        drop(file);
250    }
251
252    #[cfg(feature = "spool")]
253    #[test]
254    fn spooled_reader_supports_unknown_size_non_seekable_source() {
255        const SIZE: u64 = 1024 * 1024 + 7;
256
257        let mut source = SourceReader::streaming(io::repeat(0x5A).take(SIZE));
258
259        assert_eq!(source.try_size().unwrap(), None);
260        assert!(matches!(
261            source.seek(SeekFrom::Start(0)),
262            Err(SourceReaderError::NotSeekable)
263        ));
264
265        let mut reader = SourceReader::spooled(source).unwrap();
266
267        assert_eq!(reader.try_size().unwrap(), Some(SIZE));
268
269        assert_eq!(reader.seek(SeekFrom::End(-1)).unwrap(), SIZE - 1);
270
271        let mut buf = [0u8; 1];
272        reader.read_exact(&mut buf).unwrap();
273        assert_eq!(buf[0], 0x5A);
274
275        assert_eq!(reader.read(&mut buf).unwrap(), 0);
276
277        assert_eq!(reader.seek(SeekFrom::Start(0)).unwrap(), 0);
278        reader.read_exact(&mut buf).unwrap();
279        assert_eq!(buf[0], 0x5A);
280    }
281}
virtual_file_transfer/source_reader.rs

virtual_file_transfer/
source_reader.rs
