tonic/transport/server/service/
io.rs

1use crate::transport::server::Connected;
2use std::io;
3use std::io::IoSlice;
4use std::pin::Pin;
5use std::task::{Context, Poll};
6use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
7#[cfg(feature = "_tls-any")]
8use tokio_rustls::server::TlsStream;
9use tower_layer::Layer;
10use tower_service::Service;
11
12#[derive(Debug, Clone)]
13pub(crate) struct ConnectInfoLayer<T> {
14    connect_info: T,
15}
16
17impl<T> ConnectInfoLayer<T> {
18    pub(crate) fn new(connect_info: T) -> Self {
19        Self { connect_info }
20    }
21}
22
23impl<S, T> Layer<S> for ConnectInfoLayer<T>
24where
25    T: Clone,
26{
27    type Service = ConnectInfo<S, T>;
28
29    fn layer(&self, inner: S) -> Self::Service {
30        ConnectInfo::new(inner, self.connect_info.clone())
31    }
32}
33
34#[derive(Debug, Clone)]
35pub(crate) struct ConnectInfo<S, T> {
36    inner: S,
37    connect_info: T,
38}
39
40impl<S, T> ConnectInfo<S, T> {
41    fn new(inner: S, connect_info: T) -> Self {
42        Self {
43            inner,
44            connect_info,
45        }
46    }
47}
48
49impl<S, IO, ReqBody> Service<http::Request<ReqBody>> for ConnectInfo<S, ServerIoConnectInfo<IO>>
50where
51    S: Service<http::Request<ReqBody>>,
52    IO: Connected,
53{
54    type Response = S::Response;
55    type Error = S::Error;
56    type Future = S::Future;
57
58    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
59        self.inner.poll_ready(cx)
60    }
61
62    fn call(&mut self, mut req: http::Request<ReqBody>) -> Self::Future {
63        match self.connect_info.clone() {
64            ServerIoConnectInfo::Io(inner) => {
65                req.extensions_mut().insert(inner);
66            }
67            #[cfg(feature = "_tls-any")]
68            ServerIoConnectInfo::TlsIo(inner) => {
69                req.extensions_mut().insert(inner.get_ref().clone());
70                req.extensions_mut().insert(inner);
71            }
72        }
73        self.inner.call(req)
74    }
75}
76
77pub(crate) enum ServerIo<IO> {
78    Io(IO),
79    #[cfg(feature = "_tls-any")]
80    TlsIo(Box<TlsStream<IO>>),
81}
82
83pub(crate) enum ServerIoConnectInfo<IO: Connected> {
84    Io(<IO as Connected>::ConnectInfo),
85    #[cfg(feature = "_tls-any")]
86    TlsIo(<TlsStream<IO> as Connected>::ConnectInfo),
87}
88
89impl<IO: Connected> Clone for ServerIoConnectInfo<IO> {
90    fn clone(&self) -> Self {
91        match self {
92            Self::Io(io) => Self::Io(io.clone()),
93            #[cfg(feature = "_tls-any")]
94            Self::TlsIo(io) => Self::TlsIo(io.clone()),
95        }
96    }
97}
98
99impl<IO> ServerIo<IO> {
100    pub(in crate::transport) fn new_io(io: IO) -> Self {
101        Self::Io(io)
102    }
103
104    #[cfg(feature = "_tls-any")]
105    pub(in crate::transport) fn new_tls_io(io: TlsStream<IO>) -> Self {
106        Self::TlsIo(Box::new(io))
107    }
108
109    pub(in crate::transport) fn connect_info(&self) -> ServerIoConnectInfo<IO>
110    where
111        IO: Connected,
112    {
113        match self {
114            Self::Io(io) => ServerIoConnectInfo::Io(io.connect_info()),
115            #[cfg(feature = "_tls-any")]
116            Self::TlsIo(io) => ServerIoConnectInfo::TlsIo(io.connect_info()),
117        }
118    }
119}
120
121impl<IO> AsyncRead for ServerIo<IO>
122where
123    IO: AsyncWrite + AsyncRead + Unpin,
124{
125    fn poll_read(
126        mut self: Pin<&mut Self>,
127        cx: &mut Context<'_>,
128        buf: &mut ReadBuf<'_>,
129    ) -> Poll<io::Result<()>> {
130        match &mut *self {
131            Self::Io(io) => Pin::new(io).poll_read(cx, buf),
132            #[cfg(feature = "_tls-any")]
133            Self::TlsIo(io) => Pin::new(io).poll_read(cx, buf),
134        }
135    }
136}
137
138impl<IO> AsyncWrite for ServerIo<IO>
139where
140    IO: AsyncWrite + AsyncRead + Unpin,
141{
142    fn poll_write(
143        mut self: Pin<&mut Self>,
144        cx: &mut Context<'_>,
145        buf: &[u8],
146    ) -> Poll<io::Result<usize>> {
147        match &mut *self {
148            Self::Io(io) => Pin::new(io).poll_write(cx, buf),
149            #[cfg(feature = "_tls-any")]
150            Self::TlsIo(io) => Pin::new(io).poll_write(cx, buf),
151        }
152    }
153
154    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
155        match &mut *self {
156            Self::Io(io) => Pin::new(io).poll_flush(cx),
157            #[cfg(feature = "_tls-any")]
158            Self::TlsIo(io) => Pin::new(io).poll_flush(cx),
159        }
160    }
161
162    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
163        match &mut *self {
164            Self::Io(io) => Pin::new(io).poll_shutdown(cx),
165            #[cfg(feature = "_tls-any")]
166            Self::TlsIo(io) => Pin::new(io).poll_shutdown(cx),
167        }
168    }
169
170    fn poll_write_vectored(
171        mut self: Pin<&mut Self>,
172        cx: &mut Context<'_>,
173        bufs: &[IoSlice<'_>],
174    ) -> Poll<Result<usize, io::Error>> {
175        match &mut *self {
176            Self::Io(io) => Pin::new(io).poll_write_vectored(cx, bufs),
177            #[cfg(feature = "_tls-any")]
178            Self::TlsIo(io) => Pin::new(io).poll_write_vectored(cx, bufs),
179        }
180    }
181
182    fn is_write_vectored(&self) -> bool {
183        match self {
184            Self::Io(io) => io.is_write_vectored(),
185            #[cfg(feature = "_tls-any")]
186            Self::TlsIo(io) => io.is_write_vectored(),
187        }
188    }
189}
tonic/transport/server/service/io.rs

tonic/transport/server/service/
io.rs
