tonic/transport/service/
io.rs

1use crate::transport::server::Connected;
2use hyper::client::connect::{Connected as HyperConnected, Connection};
3use std::io;
4use std::pin::Pin;
5use std::task::{Context, Poll};
6use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
7#[cfg(feature = "tls")]
8use tokio_rustls::server::TlsStream;
9
10pub(in crate::transport) trait Io:
11    AsyncRead + AsyncWrite + Send + 'static
12{
13}
14
15impl<T> Io for T where T: AsyncRead + AsyncWrite + Send + 'static {}
16
17pub(crate) struct BoxedIo(Pin<Box<dyn Io>>);
18
19impl BoxedIo {
20    pub(in crate::transport) fn new<I: Io>(io: I) -> Self {
21        BoxedIo(Box::pin(io))
22    }
23}
24
25impl Connection for BoxedIo {
26    fn connected(&self) -> HyperConnected {
27        HyperConnected::new()
28    }
29}
30
31impl Connected for BoxedIo {
32    type ConnectInfo = NoneConnectInfo;
33
34    fn connect_info(&self) -> Self::ConnectInfo {
35        NoneConnectInfo
36    }
37}
38
39#[derive(Copy, Clone)]
40pub(crate) struct NoneConnectInfo;
41
42impl AsyncRead for BoxedIo {
43    fn poll_read(
44        mut self: Pin<&mut Self>,
45        cx: &mut Context<'_>,
46        buf: &mut ReadBuf<'_>,
47    ) -> Poll<io::Result<()>> {
48        Pin::new(&mut self.0).poll_read(cx, buf)
49    }
50}
51
52impl AsyncWrite for BoxedIo {
53    fn poll_write(
54        mut self: Pin<&mut Self>,
55        cx: &mut Context<'_>,
56        buf: &[u8],
57    ) -> Poll<io::Result<usize>> {
58        Pin::new(&mut self.0).poll_write(cx, buf)
59    }
60
61    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
62        Pin::new(&mut self.0).poll_flush(cx)
63    }
64
65    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
66        Pin::new(&mut self.0).poll_shutdown(cx)
67    }
68}
69
70pub(crate) enum ServerIo<IO> {
71    Io(IO),
72    #[cfg(feature = "tls")]
73    TlsIo(TlsStream<IO>),
74}
75
76use tower::util::Either;
77
78#[cfg(feature = "tls")]
79type ServerIoConnectInfo<IO> =
80    Either<<IO as Connected>::ConnectInfo, <TlsStream<IO> as Connected>::ConnectInfo>;
81
82#[cfg(not(feature = "tls"))]
83type ServerIoConnectInfo<IO> = Either<<IO as Connected>::ConnectInfo, ()>;
84
85impl<IO> ServerIo<IO> {
86    pub(in crate::transport) fn new_io(io: IO) -> Self {
87        Self::Io(io)
88    }
89
90    #[cfg(feature = "tls")]
91    pub(in crate::transport) fn new_tls_io(io: TlsStream<IO>) -> Self {
92        Self::TlsIo(io)
93    }
94
95    #[cfg(feature = "tls")]
96    pub(in crate::transport) fn connect_info(&self) -> ServerIoConnectInfo<IO>
97    where
98        IO: Connected,
99        TlsStream<IO>: Connected,
100    {
101        match self {
102            Self::Io(io) => Either::A(io.connect_info()),
103            Self::TlsIo(io) => Either::B(io.connect_info()),
104        }
105    }
106
107    #[cfg(not(feature = "tls"))]
108    pub(in crate::transport) fn connect_info(&self) -> ServerIoConnectInfo<IO>
109    where
110        IO: Connected,
111    {
112        match self {
113            Self::Io(io) => Either::A(io.connect_info()),
114        }
115    }
116}
117
118impl<IO> AsyncRead for ServerIo<IO>
119where
120    IO: AsyncWrite + AsyncRead + Unpin,
121{
122    fn poll_read(
123        mut self: Pin<&mut Self>,
124        cx: &mut Context<'_>,
125        buf: &mut ReadBuf<'_>,
126    ) -> Poll<io::Result<()>> {
127        match &mut *self {
128            Self::Io(io) => Pin::new(io).poll_read(cx, buf),
129            #[cfg(feature = "tls")]
130            Self::TlsIo(io) => Pin::new(io).poll_read(cx, buf),
131        }
132    }
133}
134
135impl<IO> AsyncWrite for ServerIo<IO>
136where
137    IO: AsyncWrite + AsyncRead + Unpin,
138{
139    fn poll_write(
140        mut self: Pin<&mut Self>,
141        cx: &mut Context<'_>,
142        buf: &[u8],
143    ) -> Poll<io::Result<usize>> {
144        match &mut *self {
145            Self::Io(io) => Pin::new(io).poll_write(cx, buf),
146            #[cfg(feature = "tls")]
147            Self::TlsIo(io) => Pin::new(io).poll_write(cx, buf),
148        }
149    }
150
151    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
152        match &mut *self {
153            Self::Io(io) => Pin::new(io).poll_flush(cx),
154            #[cfg(feature = "tls")]
155            Self::TlsIo(io) => Pin::new(io).poll_flush(cx),
156        }
157    }
158
159    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
160        match &mut *self {
161            Self::Io(io) => Pin::new(io).poll_shutdown(cx),
162            #[cfg(feature = "tls")]
163            Self::TlsIo(io) => Pin::new(io).poll_shutdown(cx),
164        }
165    }
166}
tonic/transport/service/io.rs

tonic/transport/service/
io.rs
