1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
use futures::Stream;
use std::path::PathBuf;
use tokio::fs::File;
use tokio_util::codec::{BytesCodec, FramedRead};

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

use std::io::Error;

use bytes::Bytes;

// Convenience wrapper around streaming out files.  Requires tokio
pub struct FileStream {
    inner: Option<FramedRead<File, BytesCodec>>,
    file: Pin<Box<dyn Future<Output = Result<File, Error>> + Send + Sync>>,
}

impl FileStream {
    pub fn new<P: Into<PathBuf>>(file: P) -> Self {
        FileStream {
            file: Box::pin(File::open(file.into())),
            inner: None,
        }
    }
}

impl Stream for FileStream {
    type Item = Result<Bytes, Error>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        if let Some(ref mut stream) = self.inner {
            return Pin::new(stream)
                .poll_next(cx)
                .map(|val| val.map(|val| val.map(|val| val.freeze())));
        } else {
            if let Poll::Ready(file_result) = self.file.as_mut().poll(cx) {
                match file_result {
                    Ok(file) => {
                        self.inner = Some(FramedRead::new(file, BytesCodec::new()));
                        cx.waker().wake_by_ref();
                    }
                    Err(err) => {
                        return Poll::Ready(Some(Err(err)));
                    }
                }
            }
        }

        Poll::Pending
    }
}

#[cfg(test)]
mod tests {
    use super::FileStream;
    use bytes::Bytes;
    use std::io::Error;
    use tokio::stream::StreamExt;

    #[tokio::test]
    async fn streams_file() -> Result<(), Error> {
        let bytes = FileStream::new("Cargo.toml")
            .collect::<Result<Bytes, Error>>()
            .await?;

        assert_eq!(bytes, &include_bytes!("../Cargo.toml")[..]);

        Ok(())
    }
}