1 use std::error::Error; 2 use std::pin::Pin; 3 use std::task::{Context, Poll}; 4 5 use bytes::{Buf, BufMut, Bytes, BytesMut}; 6 use futures_core::{ready, Stream}; 7 use http::{header, HeaderMap, HeaderValue}; 8 use http_body::{Body, SizeHint}; 9 use pin_project::pin_project; 10 use tonic::Status; 11 12 use self::content_types::*; 13 14 pub(crate) mod content_types { 15 use http::{header::CONTENT_TYPE, HeaderMap}; 16 17 pub(crate) const GRPC_WEB: &str = "application/grpc-web"; 18 pub(crate) const GRPC_WEB_PROTO: &str = "application/grpc-web+proto"; 19 pub(crate) const GRPC_WEB_TEXT: &str = "application/grpc-web-text"; 20 pub(crate) const GRPC_WEB_TEXT_PROTO: &str = "application/grpc-web-text+proto"; 21 22 pub(crate) fn is_grpc_web(headers: &HeaderMap) -> bool { 23 matches!( 24 content_type(headers), 25 Some(GRPC_WEB) | Some(GRPC_WEB_PROTO) | Some(GRPC_WEB_TEXT) | Some(GRPC_WEB_TEXT_PROTO) 26 ) 27 } 28 29 fn content_type(headers: &HeaderMap) -> Option<&str> { 30 headers.get(CONTENT_TYPE).and_then(|val| val.to_str().ok()) 31 } 32 } 33 34 const BUFFER_SIZE: usize = 8 * 1024; 35 36 const FRAME_HEADER_SIZE: usize = 5; 37 38 // 8th (MSB) bit of the 1st gRPC frame byte 39 // denotes an uncompressed trailer (as part of the body) 40 const GRPC_WEB_TRAILERS_BIT: u8 = 0b10000000; 41 42 #[derive(Copy, Clone, PartialEq, Debug)] 43 enum Direction { 44 Request, 45 Response, 46 } 47 48 #[derive(Copy, Clone, PartialEq, Debug)] 49 pub(crate) enum Encoding { 50 Base64, 51 None, 52 } 53 54 #[pin_project] 55 pub(crate) struct GrpcWebCall<B> { 56 #[pin] 57 inner: B, 58 buf: BytesMut, 59 direction: Direction, 60 encoding: Encoding, 61 poll_trailers: bool, 62 } 63 64 impl<B> GrpcWebCall<B> { 65 pub(crate) fn request(inner: B, encoding: Encoding) -> Self { 66 Self::new(inner, Direction::Request, encoding) 67 } 68 69 pub(crate) fn response(inner: B, encoding: Encoding) -> Self { 70 Self::new(inner, Direction::Response, encoding) 71 } 72 73 fn new(inner: B, direction: Direction, encoding: Encoding) -> Self { 74 GrpcWebCall { 75 inner, 76 buf: BytesMut::with_capacity(match (direction, encoding) { 77 (Direction::Response, Encoding::Base64) => BUFFER_SIZE, 78 _ => 0, 79 }), 80 direction, 81 encoding, 82 poll_trailers: true, 83 } 84 } 85 86 // This is to avoid passing a slice of bytes with a length that the base64 87 // decoder would consider invalid. 88 #[inline] 89 fn max_decodable(&self) -> usize { 90 (self.buf.len() / 4) * 4 91 } 92 93 fn decode_chunk(mut self: Pin<&mut Self>) -> Result<Option<Bytes>, Status> { 94 // not enough bytes to decode 95 if self.buf.is_empty() || self.buf.len() < 4 { 96 return Ok(None); 97 } 98 99 // Split `buf` at the largest index that is multiple of 4. Decode the 100 // returned `Bytes`, keeping the rest for the next attempt to decode. 101 let index = self.max_decodable(); 102 103 base64::decode(self.as_mut().project().buf.split_to(index)) 104 .map(|decoded| Some(Bytes::from(decoded))) 105 .map_err(internal_error) 106 } 107 } 108 109 impl<B> GrpcWebCall<B> 110 where 111 B: Body<Data = Bytes>, 112 B::Error: Error, 113 { 114 fn poll_decode( 115 mut self: Pin<&mut Self>, 116 cx: &mut Context<'_>, 117 ) -> Poll<Option<Result<B::Data, Status>>> { 118 match self.encoding { 119 Encoding::Base64 => loop { 120 if let Some(bytes) = self.as_mut().decode_chunk()? { 121 return Poll::Ready(Some(Ok(bytes))); 122 } 123 124 let mut this = self.as_mut().project(); 125 126 match ready!(this.inner.as_mut().poll_data(cx)) { 127 Some(Ok(data)) => this.buf.put(data), 128 Some(Err(e)) => return Poll::Ready(Some(Err(internal_error(e)))), 129 None => { 130 return if this.buf.has_remaining() { 131 Poll::Ready(Some(Err(internal_error("malformed base64 request")))) 132 } else { 133 Poll::Ready(None) 134 } 135 } 136 } 137 }, 138 139 Encoding::None => match ready!(self.project().inner.poll_data(cx)) { 140 Some(res) => Poll::Ready(Some(res.map_err(internal_error))), 141 None => Poll::Ready(None), 142 }, 143 } 144 } 145 146 fn poll_encode( 147 mut self: Pin<&mut Self>, 148 cx: &mut Context<'_>, 149 ) -> Poll<Option<Result<B::Data, Status>>> { 150 let mut this = self.as_mut().project(); 151 152 if let Some(mut res) = ready!(this.inner.as_mut().poll_data(cx)) { 153 if *this.encoding == Encoding::Base64 { 154 res = res.map(|b| base64::encode(b).into()) 155 } 156 157 return Poll::Ready(Some(res.map_err(internal_error))); 158 } 159 160 // this flag is needed because the inner stream never 161 // returns Poll::Ready(None) when polled for trailers 162 if *this.poll_trailers { 163 return match ready!(this.inner.poll_trailers(cx)) { 164 Ok(Some(map)) => { 165 let mut frame = make_trailers_frame(map); 166 167 if *this.encoding == Encoding::Base64 { 168 frame = base64::encode(frame).into_bytes(); 169 } 170 171 *this.poll_trailers = false; 172 Poll::Ready(Some(Ok(frame.into()))) 173 } 174 Ok(None) => Poll::Ready(None), 175 Err(e) => Poll::Ready(Some(Err(internal_error(e)))), 176 }; 177 } 178 179 Poll::Ready(None) 180 } 181 } 182 183 impl<B> Body for GrpcWebCall<B> 184 where 185 B: Body<Data = Bytes>, 186 B::Error: Error, 187 { 188 type Data = Bytes; 189 type Error = Status; 190 191 fn poll_data( 192 self: Pin<&mut Self>, 193 cx: &mut Context<'_>, 194 ) -> Poll<Option<Result<Self::Data, Self::Error>>> { 195 match self.direction { 196 Direction::Request => self.poll_decode(cx), 197 Direction::Response => self.poll_encode(cx), 198 } 199 } 200 201 fn poll_trailers( 202 self: Pin<&mut Self>, 203 _: &mut Context<'_>, 204 ) -> Poll<Result<Option<HeaderMap<HeaderValue>>, Self::Error>> { 205 Poll::Ready(Ok(None)) 206 } 207 208 fn is_end_stream(&self) -> bool { 209 self.inner.is_end_stream() 210 } 211 212 fn size_hint(&self) -> SizeHint { 213 self.inner.size_hint() 214 } 215 } 216 217 impl<B> Stream for GrpcWebCall<B> 218 where 219 B: Body<Data = Bytes>, 220 B::Error: Error, 221 { 222 type Item = Result<Bytes, Status>; 223 224 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { 225 Body::poll_data(self, cx) 226 } 227 } 228 229 impl Encoding { 230 pub(crate) fn from_content_type(headers: &HeaderMap) -> Encoding { 231 Self::from_header(headers.get(header::CONTENT_TYPE)) 232 } 233 234 pub(crate) fn from_accept(headers: &HeaderMap) -> Encoding { 235 Self::from_header(headers.get(header::ACCEPT)) 236 } 237 238 pub(crate) fn to_content_type(&self) -> &'static str { 239 match self { 240 Encoding::Base64 => GRPC_WEB_TEXT_PROTO, 241 Encoding::None => GRPC_WEB_PROTO, 242 } 243 } 244 245 fn from_header(value: Option<&HeaderValue>) -> Encoding { 246 match value.and_then(|val| val.to_str().ok()) { 247 Some(GRPC_WEB_TEXT_PROTO) | Some(GRPC_WEB_TEXT) => Encoding::Base64, 248 _ => Encoding::None, 249 } 250 } 251 } 252 253 fn internal_error(e: impl std::fmt::Display) -> Status { 254 Status::internal(format!("tonic-web: {}", e)) 255 } 256 257 // Key-value pairs encoded as a HTTP/1 headers block (without the terminating newline) 258 fn encode_trailers(trailers: HeaderMap) -> Vec<u8> { 259 trailers.iter().fold(Vec::new(), |mut acc, (key, value)| { 260 acc.put_slice(key.as_ref()); 261 acc.push(b':'); 262 acc.put_slice(value.as_bytes()); 263 acc.put_slice(b"\r\n"); 264 acc 265 }) 266 } 267 268 fn make_trailers_frame(trailers: HeaderMap) -> Vec<u8> { 269 let trailers = encode_trailers(trailers); 270 let len = trailers.len(); 271 assert!(len <= u32::MAX as usize); 272 273 let mut frame = Vec::with_capacity(len + FRAME_HEADER_SIZE); 274 frame.push(GRPC_WEB_TRAILERS_BIT); 275 frame.put_u32(len as u32); 276 frame.extend(trailers); 277 278 frame 279 } 280 281 #[cfg(test)] 282 mod tests { 283 use super::*; 284 285 #[test] 286 fn encoding_constructors() { 287 let cases = &[ 288 (GRPC_WEB, Encoding::None), 289 (GRPC_WEB_PROTO, Encoding::None), 290 (GRPC_WEB_TEXT, Encoding::Base64), 291 (GRPC_WEB_TEXT_PROTO, Encoding::Base64), 292 ("foo", Encoding::None), 293 ]; 294 295 let mut headers = HeaderMap::new(); 296 297 for case in cases { 298 headers.insert(header::CONTENT_TYPE, case.0.parse().unwrap()); 299 headers.insert(header::ACCEPT, case.0.parse().unwrap()); 300 301 assert_eq!(Encoding::from_content_type(&headers), case.1, "{}", case.0); 302 assert_eq!(Encoding::from_accept(&headers), case.1, "{}", case.0); 303 } 304 } 305 } 306