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