structured_proxy/
upstream.rs1use std::convert::Infallible;
10use std::future::Future;
11use std::pin::Pin;
12use std::task::{ready, Context, Poll};
13
14use bytes::Bytes;
15use pin_project_lite::pin_project;
16use tonic::client::GrpcService;
17
18pub type BoxError = Box<dyn std::error::Error + Send + Sync>;
20
21pub trait Upstream:
41 GrpcService<
42 tonic::body::Body,
43 ResponseBody: http_body::Body<Data = Bytes, Error: Into<BoxError> + Send> + Send + 'static,
44 Error: Into<BoxError>,
45 Future: Send,
46 > + Clone
47 + Send
48 + Sync
49 + 'static
50{
51}
52
53impl<T> Upstream for T where
54 T: GrpcService<
55 tonic::body::Body,
56 ResponseBody: http_body::Body<Data = Bytes, Error: Into<BoxError> + Send>
57 + Send
58 + 'static,
59 Error: Into<BoxError>,
60 Future: Send,
61 > + Clone
62 + Send
63 + Sync
64 + 'static
65{
66}
67
68pin_project! {
69 #[project = PassThroughProj]
81 #[project_replace = PassThroughReplace]
82 pub(crate) enum PassThrough<U: Upstream> {
83 Ready {
84 upstream: U,
85 request: http::Request<tonic::body::Body>,
86 protocol: GrpcProtocol,
87 },
88 Call {
89 #[pin]
90 future: U::Future,
91 protocol: GrpcProtocol,
92 },
93 Done,
94 }
95}
96
97impl<U: Upstream> PassThrough<U> {
98 pub(crate) fn new(
99 upstream: U,
100 request: http::Request<tonic::body::Body>,
101 protocol: GrpcProtocol,
102 ) -> Self {
103 Self::Ready {
104 upstream,
105 request,
106 protocol,
107 }
108 }
109}
110
111impl<U: Upstream> Future for PassThrough<U> {
112 type Output = Result<http::Response<axum::body::Body>, Infallible>;
113
114 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
115 loop {
116 match self.as_mut().project() {
117 PassThroughProj::Ready {
118 upstream, protocol, ..
119 } => {
120 let protocol = *protocol;
121 if let Err(e) = ready!(upstream.poll_ready(cx)) {
122 self.set(Self::Done);
123 return Poll::Ready(Ok(failure(e.into(), protocol)));
124 }
125 let PassThroughReplace::Ready {
126 mut upstream,
127 request,
128 ..
129 } = self.as_mut().project_replace(Self::Done)
130 else {
131 unreachable!("the state was just matched as Ready");
132 };
133 let future = upstream.call(request);
136 self.set(Self::Call { future, protocol });
137 }
138 PassThroughProj::Call { future, protocol } => {
139 let protocol = *protocol;
140 let result = ready!(future.poll(cx));
141 self.set(Self::Done);
142 return Poll::Ready(Ok(match result {
143 Ok(response) => response.map(axum::body::Body::new),
144 Err(e) => failure(e.into(), protocol),
145 }));
146 }
147 PassThroughProj::Done => panic!("PassThrough polled after completion"),
148 }
149 }
150 }
151}
152
153fn failure(error: BoxError, protocol: GrpcProtocol) -> http::Response<axum::body::Body> {
158 trailers_only(tonic::Status::from_error(error), protocol)
159}
160
161pub(crate) fn trailers_only(
163 status: tonic::Status,
164 protocol: GrpcProtocol,
165) -> http::Response<axum::body::Body> {
166 let mut response: http::Response<axum::body::Body> = status.into_http();
167 if protocol != GrpcProtocol::Grpc {
168 response.headers_mut().insert(
169 http::header::CONTENT_TYPE,
170 http::HeaderValue::from_static(protocol.content_type()),
171 );
172 }
173 response
174}
175
176#[derive(Clone, Copy, Debug, PartialEq, Eq)]
178pub(crate) enum GrpcProtocol {
179 Grpc,
181 Web,
183 WebText,
185}
186
187impl GrpcProtocol {
188 fn content_type(self) -> &'static str {
191 match self {
192 Self::Grpc => "application/grpc",
193 Self::Web => "application/grpc-web+proto",
194 Self::WebText => "application/grpc-web-text+proto",
195 }
196 }
197}
198
199pub(crate) fn grpc_protocol(headers: &http::HeaderMap) -> Option<GrpcProtocol> {
204 let value = headers.get(http::header::CONTENT_TYPE)?.as_bytes();
205 let media = match value.iter().position(|&b| b == b';') {
206 Some(end) => &value[..end],
207 None => value,
208 }
209 .trim_ascii();
210 let rest = strip_prefix_ignore_case(media, b"application/grpc")?;
211 match rest {
212 [] | [b'+', ..] => Some(GrpcProtocol::Grpc),
213 _ => web_protocol(strip_prefix_ignore_case(rest, b"-web")?),
214 }
215}
216
217pub(crate) fn is_grpc_web_preflight(method: &http::Method, headers: &http::HeaderMap) -> bool {
224 method == http::Method::OPTIONS
225 && headers.contains_key(http::header::ORIGIN)
226 && headers.contains_key(http::header::ACCESS_CONTROL_REQUEST_METHOD)
227 && headers
228 .get_all(http::header::ACCESS_CONTROL_REQUEST_HEADERS)
229 .iter()
230 .flat_map(|value| value.as_bytes().split(|&b| b == b','))
231 .any(|name| name.trim_ascii().eq_ignore_ascii_case(b"x-grpc-web"))
232}
233
234fn web_protocol(rest: &[u8]) -> Option<GrpcProtocol> {
237 let (protocol, rest) = match strip_prefix_ignore_case(rest, b"-text") {
238 Some(rest) => (GrpcProtocol::WebText, rest),
239 None => (GrpcProtocol::Web, rest),
240 };
241 matches!(rest, [] | [b'+', ..]).then_some(protocol)
242}
243
244fn strip_prefix_ignore_case<'a>(bytes: &'a [u8], prefix: &[u8]) -> Option<&'a [u8]> {
246 let (head, rest) = bytes.split_at_checked(prefix.len())?;
247 head.eq_ignore_ascii_case(prefix).then_some(rest)
248}
249
250#[cfg(test)]
251mod tests;