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
use futures_util::stream::{Stream, StreamExt};
use hyper::body::{Body, Bytes, Sender};
use std::io::Error as IoError;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::fs::File;
use tokio::prelude::AsyncRead;

const BUF_SIZE: usize = 8 * 1024;

/// Wraps a `tokio::fs::File`, and implements a stream of `Bytes`s.
pub struct FileBytesStream {
    file: File,
    buf: Box<[u8; BUF_SIZE]>,
}

impl FileBytesStream {
    /// Create a new stream from the given file.
    pub fn new(file: File) -> FileBytesStream {
        let buf = Box::new([0; BUF_SIZE]);
        FileBytesStream { file, buf }
    }
}

impl Stream for FileBytesStream {
    type Item = Result<Bytes, IoError>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        let Self {
            ref mut file,
            ref mut buf,
        } = *self;
        match Pin::new(file).poll_read(cx, &mut buf[..]) {
            Poll::Ready(Ok(0)) => Poll::Ready(None),
            Poll::Ready(Ok(size)) => Poll::Ready(Some(Ok(self.buf[..size].to_owned().into()))),
            Poll::Ready(Err(e)) => Poll::Ready(Some(Err(e))),
            Poll::Pending => Poll::Pending,
        }
    }
}

impl FileBytesStream {
    /// Create a Hyper `Body` from this stream.
    pub fn into_body(self) -> Body {
        let (sender, body) = Body::channel();
        tokio::spawn(self.body_sender_loop(sender));
        body
    }

    async fn body_sender_loop(mut self, mut sender: Sender) {
        loop {
            let (result, stream) = self.into_future().await;
            self = stream;

            let chunk = match result {
                Some(Ok(chunk)) => chunk,
                Some(Err(_)) => return sender.abort(),
                None => break,
            };

            if let Err(_) = sender.send_data(chunk).await {
                break;
            }
        }
    }
}