rocket/data/
io_stream.rs

Help
1use std::io;
2use std::task::{Context, Poll};
3use std::pin::Pin;
4
5use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
6
7use crate::http::hyper::upgrade::Upgraded;
8
9/// A bidirectional, raw stream to the client.
10///
11/// An instance of `IoStream` is passed to an [`IoHandler`] in response to a
12/// successful upgrade request initiated by responders via
13/// [`Response::add_upgrade()`] or the equivalent builder method
14/// [`Builder::upgrade()`]. For details on upgrade connections, see
15/// [`Response`#upgrading].
16///
17/// An `IoStream` is guaranteed to be [`AsyncRead`], [`AsyncWrite`], and
18/// `Unpin`. Bytes written to the stream are sent directly to the client. Bytes
19/// read from the stream are those sent directly _by_ the client. See
20/// [`IoHandler`] for one example of how values of this type are used.
21///
22/// [`Response::add_upgrade()`]: crate::Response::add_upgrade()
23/// [`Builder::upgrade()`]: crate::response::Builder::upgrade()
24/// [`Response`#upgrading]: crate::response::Response#upgrading
25pub struct IoStream {
26    kind: IoStreamKind,
27}
28
29/// Just in case we want to add stream kinds in the future.
30enum IoStreamKind {
31    Upgraded(Upgraded)
32}
33
34/// An upgraded connection I/O handler.
35///
36/// An I/O handler performs raw I/O via the passed in [`IoStream`], which is
37/// [`AsyncRead`], [`AsyncWrite`], and `Unpin`.
38///
39/// # Example
40///
41/// The example below implements an `EchoHandler` that echos the raw bytes back
42/// to the client.
43///
44/// ```rust
45/// use std::pin::Pin;
46///
47/// use rocket::tokio::io;
48/// use rocket::data::{IoHandler, IoStream};
49///
50/// struct EchoHandler;
51///
52/// #[rocket::async_trait]
53/// impl IoHandler for EchoHandler {
54///     async fn io(self: Pin<Box<Self>>, io: IoStream) -> io::Result<()> {
55///         let (mut reader, mut writer) = io::split(io);
56///         io::copy(&mut reader, &mut writer).await?;
57///         Ok(())
58///     }
59/// }
60///
61/// # use rocket::Response;
62/// # rocket::async_test(async {
63/// # let mut response = Response::new();
64/// # response.add_upgrade("raw-echo", EchoHandler);
65/// # assert!(response.upgrade("raw-echo").is_some());
66/// # })
67/// ```
68#[crate::async_trait]
69pub trait IoHandler: Send {
70    /// Performs the raw I/O.
71    async fn io(self: Pin<Box<Self>>, io: IoStream) -> io::Result<()>;
72}
73
74#[doc(hidden)]
75impl From<Upgraded> for IoStream {
76    fn from(io: Upgraded) -> Self {
77        IoStream { kind: IoStreamKind::Upgraded(io) }
78    }
79}
80
81/// A "trait alias" of sorts so we can use `AsyncRead + AsyncWrite + Unpin` in `dyn`.
82pub trait AsyncReadWrite: AsyncRead + AsyncWrite + Unpin { }
83
84/// Implemented for all `AsyncRead + AsyncWrite + Unpin`, of course.
85impl<T: AsyncRead + AsyncWrite + Unpin> AsyncReadWrite for T {  }
86
87impl IoStream {
88    /// Returns the internal I/O stream.
89    fn inner_mut(&mut self) -> Pin<&mut dyn AsyncReadWrite> {
90        match self.kind {
91            IoStreamKind::Upgraded(ref mut io) => Pin::new(io),
92        }
93    }
94
95    /// Returns `true` if the inner I/O stream is write vectored.
96    fn inner_is_write_vectored(&self) -> bool {
97        match self.kind {
98            IoStreamKind::Upgraded(ref io) => io.is_write_vectored(),
99        }
100    }
101}
102
103impl AsyncRead for IoStream {
104    fn poll_read(
105        self: Pin<&mut Self>,
106        cx: &mut Context<'_>,
107        buf: &mut ReadBuf<'_>,
108    ) -> Poll<io::Result<()>> {
109        self.get_mut().inner_mut().poll_read(cx, buf)
110    }
111}
112
113impl AsyncWrite for IoStream {
114    fn poll_write(
115        self: Pin<&mut Self>,
116        cx: &mut Context<'_>,
117        buf: &[u8],
118    ) -> Poll<io::Result<usize>> {
119        self.get_mut().inner_mut().poll_write(cx, buf)
120    }
121
122    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
123        self.get_mut().inner_mut().poll_flush(cx)
124    }
125
126    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
127        self.get_mut().inner_mut().poll_shutdown(cx)
128    }
129
130    fn poll_write_vectored(
131        self: Pin<&mut Self>,
132        cx: &mut Context<'_>,
133        bufs: &[io::IoSlice<'_>],
134    ) -> Poll<io::Result<usize>> {
135        self.get_mut().inner_mut().poll_write_vectored(cx, bufs)
136    }
137
138    fn is_write_vectored(&self) -> bool {
139        self.inner_is_write_vectored()
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[test]
148    fn is_unpin() {
149        fn check_traits<T: AsyncRead + AsyncWrite + Unpin + Send>() {}
150        check_traits::<IoStream>();
151    }
152}
rocket/data/io_stream.rs

rocket/data/
io_stream.rs
