1 use std::error::Error; 2 use std::pin::Pin; 3 use std::task::{ready, Context, Poll}; 4 5 use base64::Engine as _; 6 use bytes::{Buf, BufMut, Bytes, BytesMut}; 7 use http::{header, HeaderMap, HeaderName, HeaderValue}; 8 use http_body::{Body, SizeHint}; 9 use pin_project::pin_project; 10 use tokio_stream::Stream; 11 use tonic::Status; 12 13 use self::content_types::*; 14 15 // A grpc header is u8 (flag) + u32 (msg len) 16 const GRPC_HEADER_SIZE: usize = 1 + 4; 17 18 pub(crate) mod content_types { 19 use http::{header::CONTENT_TYPE, HeaderMap}; 20 21 pub(crate) const GRPC_WEB: &str = "application/grpc-web"; 22 pub(crate) const GRPC_WEB_PROTO: &str = "application/grpc-web+proto"; 23 pub(crate) const GRPC_WEB_TEXT: &str = "application/grpc-web-text"; 24 pub(crate) const GRPC_WEB_TEXT_PROTO: &str = "application/grpc-web-text+proto"; 25 26 pub(crate) fn is_grpc_web(headers: &HeaderMap) -> bool { 27 matches!( 28 content_type(headers), 29 Some(GRPC_WEB) | Some(GRPC_WEB_PROTO) | Some(GRPC_WEB_TEXT) | Some(GRPC_WEB_TEXT_PROTO) 30 ) 31 } 32 33 fn content_type(headers: &HeaderMap) -> Option<&str> { 34 headers.get(CONTENT_TYPE).and_then(|val| val.to_str().ok()) 35 } 36 } 37 38 const BUFFER_SIZE: usize = 8 * 1024; 39 40 const FRAME_HEADER_SIZE: usize = 5; 41 42 // 8th (MSB) bit of the 1st gRPC frame byte 43 // denotes an uncompressed trailer (as part of the body) 44 const GRPC_WEB_TRAILERS_BIT: u8 = 0b10000000; 45 46 #[derive(Copy, Clone, PartialEq, Debug)] 47 enum Direction { 48 Decode, 49 Encode, 50 Empty, 51 } 52 53 #[derive(Copy, Clone, PartialEq, Debug)] 54 pub(crate) enum Encoding { 55 Base64, 56 None, 57 } 58 59 /// HttpBody adapter for the grpc web based services. 60 #[derive(Debug)] 61 #[pin_project] 62 pub struct GrpcWebCall<B> { 63 #[pin] 64 inner: B, 65 buf: BytesMut, 66 direction: Direction, 67 encoding: Encoding, 68 poll_trailers: bool, 69 client: bool, 70 trailers: Option<HeaderMap>, 71 } 72 73 impl<B: Default> Default for GrpcWebCall<B> { 74 fn default() -> Self { 75 Self { 76 inner: Default::default(), 77 buf: Default::default(), 78 direction: Direction::Empty, 79 encoding: Encoding::None, 80 poll_trailers: Default::default(), 81 client: Default::default(), 82 trailers: Default::default(), 83 } 84 } 85 } 86 87 impl<B> GrpcWebCall<B> { 88 pub(crate) fn request(inner: B, encoding: Encoding) -> Self { 89 Self::new(inner, Direction::Decode, encoding) 90 } 91 92 pub(crate) fn response(inner: B, encoding: Encoding) -> Self { 93 Self::new(inner, Direction::Encode, encoding) 94 } 95 96 pub(crate) fn client_request(inner: B) -> Self { 97 Self::new_client(inner, Direction::Encode, Encoding::None) 98 } 99 100 pub(crate) fn client_response(inner: B) -> Self { 101 Self::new_client(inner, Direction::Decode, Encoding::None) 102 } 103 104 fn new_client(inner: B, direction: Direction, encoding: Encoding) -> Self { 105 GrpcWebCall { 106 inner, 107 buf: BytesMut::with_capacity(match (direction, encoding) { 108 (Direction::Encode, Encoding::Base64) => BUFFER_SIZE, 109 _ => 0, 110 }), 111 direction, 112 encoding, 113 poll_trailers: true, 114 client: true, 115 trailers: None, 116 } 117 } 118 119 fn new(inner: B, direction: Direction, encoding: Encoding) -> Self { 120 GrpcWebCall { 121 inner, 122 buf: BytesMut::with_capacity(match (direction, encoding) { 123 (Direction::Encode, Encoding::Base64) => BUFFER_SIZE, 124 _ => 0, 125 }), 126 direction, 127 encoding, 128 poll_trailers: true, 129 client: false, 130 trailers: None, 131 } 132 } 133 134 // This is to avoid passing a slice of bytes with a length that the base64 135 // decoder would consider invalid. 136 #[inline] 137 fn max_decodable(&self) -> usize { 138 (self.buf.len() / 4) * 4 139 } 140 141 fn decode_chunk(mut self: Pin<&mut Self>) -> Result<Option<Bytes>, Status> { 142 // not enough bytes to decode 143 if self.buf.is_empty() || self.buf.len() < 4 { 144 return Ok(None); 145 } 146 147 // Split `buf` at the largest index that is multiple of 4. Decode the 148 // returned `Bytes`, keeping the rest for the next attempt to decode. 149 let index = self.max_decodable(); 150 151 crate::util::base64::STANDARD 152 .decode(self.as_mut().project().buf.split_to(index)) 153 .map(|decoded| Some(Bytes::from(decoded))) 154 .map_err(internal_error) 155 } 156 } 157 158 impl<B> GrpcWebCall<B> 159 where 160 B: Body<Data = Bytes>, 161 B::Error: Error, 162 { 163 fn poll_decode( 164 mut self: Pin<&mut Self>, 165 cx: &mut Context<'_>, 166 ) -> Poll<Option<Result<B::Data, Status>>> { 167 match self.encoding { 168 Encoding::Base64 => loop { 169 if let Some(bytes) = self.as_mut().decode_chunk()? { 170 return Poll::Ready(Some(Ok(bytes))); 171 } 172 173 let mut this = self.as_mut().project(); 174 175 match ready!(this.inner.as_mut().poll_data(cx)) { 176 Some(Ok(data)) => this.buf.put(data), 177 Some(Err(e)) => return Poll::Ready(Some(Err(internal_error(e)))), 178 None => { 179 return if this.buf.has_remaining() { 180 Poll::Ready(Some(Err(internal_error("malformed base64 request")))) 181 } else { 182 Poll::Ready(None) 183 } 184 } 185 } 186 }, 187 188 Encoding::None => match ready!(self.project().inner.poll_data(cx)) { 189 Some(res) => Poll::Ready(Some(res.map_err(internal_error))), 190 None => Poll::Ready(None), 191 }, 192 } 193 } 194 195 fn poll_encode( 196 mut self: Pin<&mut Self>, 197 cx: &mut Context<'_>, 198 ) -> Poll<Option<Result<B::Data, Status>>> { 199 let mut this = self.as_mut().project(); 200 201 if let Some(mut res) = ready!(this.inner.as_mut().poll_data(cx)) { 202 if *this.encoding == Encoding::Base64 { 203 res = res.map(|b| crate::util::base64::STANDARD.encode(b).into()) 204 } 205 206 return Poll::Ready(Some(res.map_err(internal_error))); 207 } 208 209 // this flag is needed because the inner stream never 210 // returns Poll::Ready(None) when polled for trailers 211 if *this.poll_trailers { 212 return match ready!(this.inner.poll_trailers(cx)) { 213 Ok(Some(map)) => { 214 let mut frame = make_trailers_frame(map); 215 216 if *this.encoding == Encoding::Base64 { 217 frame = crate::util::base64::STANDARD.encode(frame).into_bytes(); 218 } 219 220 *this.poll_trailers = false; 221 Poll::Ready(Some(Ok(frame.into()))) 222 } 223 Ok(None) => Poll::Ready(None), 224 Err(e) => Poll::Ready(Some(Err(internal_error(e)))), 225 }; 226 } 227 228 Poll::Ready(None) 229 } 230 } 231 232 impl<B> Body for GrpcWebCall<B> 233 where 234 B: Body<Data = Bytes>, 235 B::Error: Error, 236 { 237 type Data = Bytes; 238 type Error = Status; 239 240 fn poll_data( 241 mut self: Pin<&mut Self>, 242 cx: &mut Context<'_>, 243 ) -> Poll<Option<Result<Self::Data, Self::Error>>> { 244 if self.client && self.direction == Direction::Decode { 245 let mut me = self.as_mut(); 246 247 loop { 248 let incoming_buf = match ready!(me.as_mut().poll_decode(cx)) { 249 Some(Ok(incoming_buf)) => incoming_buf, 250 None => { 251 // TODO: Consider eofing here? 252 // Even if the buffer has more data, this will hit the eof branch 253 // of decode in tonic 254 return Poll::Ready(None); 255 } 256 Some(Err(e)) => return Poll::Ready(Some(Err(e))), 257 }; 258 259 let buf = &mut me.as_mut().project().buf; 260 261 buf.put(incoming_buf); 262 263 return match find_trailers(&buf[..])? { 264 FindTrailers::Trailer(len) => { 265 // Extract up to len of where the trailers are at 266 let msg_buf = buf.copy_to_bytes(len); 267 match decode_trailers_frame(buf.split().freeze()) { 268 Ok(Some(trailers)) => { 269 self.project().trailers.replace(trailers); 270 } 271 Err(e) => return Poll::Ready(Some(Err(e))), 272 _ => {} 273 } 274 275 if msg_buf.has_remaining() { 276 Poll::Ready(Some(Ok(msg_buf))) 277 } else { 278 Poll::Ready(None) 279 } 280 } 281 FindTrailers::IncompleteBuf => continue, 282 FindTrailers::Done(len) => Poll::Ready(Some(Ok(buf.split_to(len).freeze()))), 283 }; 284 } 285 } 286 287 match self.direction { 288 Direction::Decode => self.poll_decode(cx), 289 Direction::Encode => self.poll_encode(cx), 290 Direction::Empty => Poll::Ready(None), 291 } 292 } 293 294 fn poll_trailers( 295 self: Pin<&mut Self>, 296 _: &mut Context<'_>, 297 ) -> Poll<Result<Option<HeaderMap<HeaderValue>>, Self::Error>> { 298 let trailers = self.project().trailers.take(); 299 Poll::Ready(Ok(trailers)) 300 } 301 302 fn is_end_stream(&self) -> bool { 303 self.inner.is_end_stream() 304 } 305 306 fn size_hint(&self) -> SizeHint { 307 self.inner.size_hint() 308 } 309 } 310 311 impl<B> Stream for GrpcWebCall<B> 312 where 313 B: Body<Data = Bytes>, 314 B::Error: Error, 315 { 316 type Item = Result<Bytes, Status>; 317 318 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { 319 Body::poll_data(self, cx) 320 } 321 } 322 323 impl Encoding { 324 pub(crate) fn from_content_type(headers: &HeaderMap) -> Encoding { 325 Self::from_header(headers.get(header::CONTENT_TYPE)) 326 } 327 328 pub(crate) fn from_accept(headers: &HeaderMap) -> Encoding { 329 Self::from_header(headers.get(header::ACCEPT)) 330 } 331 332 pub(crate) fn to_content_type(self) -> &'static str { 333 match self { 334 Encoding::Base64 => GRPC_WEB_TEXT_PROTO, 335 Encoding::None => GRPC_WEB_PROTO, 336 } 337 } 338 339 fn from_header(value: Option<&HeaderValue>) -> Encoding { 340 match value.and_then(|val| val.to_str().ok()) { 341 Some(GRPC_WEB_TEXT_PROTO) | Some(GRPC_WEB_TEXT) => Encoding::Base64, 342 _ => Encoding::None, 343 } 344 } 345 } 346 347 fn internal_error(e: impl std::fmt::Display) -> Status { 348 Status::internal(format!("tonic-web: {}", e)) 349 } 350 351 // Key-value pairs encoded as a HTTP/1 headers block (without the terminating newline) 352 fn encode_trailers(trailers: HeaderMap) -> Vec<u8> { 353 trailers.iter().fold(Vec::new(), |mut acc, (key, value)| { 354 acc.put_slice(key.as_ref()); 355 acc.push(b':'); 356 acc.put_slice(value.as_bytes()); 357 acc.put_slice(b"\r\n"); 358 acc 359 }) 360 } 361 362 fn decode_trailers_frame(mut buf: Bytes) -> Result<Option<HeaderMap>, Status> { 363 if buf.remaining() < GRPC_HEADER_SIZE { 364 return Ok(None); 365 } 366 367 buf.get_u8(); 368 buf.get_u32(); 369 370 let mut map = HeaderMap::new(); 371 let mut temp_buf = buf.clone(); 372 373 let mut trailers = Vec::new(); 374 let mut cursor_pos = 0; 375 376 for (i, b) in buf.iter().enumerate() { 377 if b == &b'\r' && buf.get(i + 1) == Some(&b'\n') { 378 let trailer = temp_buf.copy_to_bytes(i - cursor_pos); 379 cursor_pos = i; 380 trailers.push(trailer); 381 if temp_buf.has_remaining() { 382 temp_buf.get_u8(); 383 temp_buf.get_u8(); 384 } 385 } 386 } 387 388 for trailer in trailers { 389 let mut s = trailer.split(|b| b == &b':'); 390 let key = s 391 .next() 392 .ok_or_else(|| Status::internal("trailers couldn't parse key"))?; 393 let value = s 394 .next() 395 .ok_or_else(|| Status::internal("trailers couldn't parse value"))?; 396 397 let value = value 398 .split(|b| b == &b'\r') 399 .next() 400 .ok_or_else(|| Status::internal("trailers was not escaped"))?; 401 402 let header_key = HeaderName::try_from(key) 403 .map_err(|e| Status::internal(format!("Unable to parse HeaderName: {}", e)))?; 404 let header_value = HeaderValue::try_from(value) 405 .map_err(|e| Status::internal(format!("Unable to parse HeaderValue: {}", e)))?; 406 map.insert(header_key, header_value); 407 } 408 409 Ok(Some(map)) 410 } 411 412 fn make_trailers_frame(trailers: HeaderMap) -> Vec<u8> { 413 let trailers = encode_trailers(trailers); 414 let len = trailers.len(); 415 assert!(len <= u32::MAX as usize); 416 417 let mut frame = Vec::with_capacity(len + FRAME_HEADER_SIZE); 418 frame.push(GRPC_WEB_TRAILERS_BIT); 419 frame.put_u32(len as u32); 420 frame.extend(trailers); 421 422 frame 423 } 424 425 /// Search some buffer for grpc-web trailers headers and return 426 /// its location in the original buf. If `None` is returned we did 427 /// not find a trailers in this buffer either because its incomplete 428 /// or the buffer just contained grpc message frames. 429 fn find_trailers(buf: &[u8]) -> Result<FindTrailers, Status> { 430 let mut len = 0; 431 let mut temp_buf = buf; 432 433 loop { 434 // To check each frame, there must be at least GRPC_HEADER_SIZE 435 // amount of bytes available otherwise the buffer is incomplete. 436 if temp_buf.is_empty() || temp_buf.len() < GRPC_HEADER_SIZE { 437 return Ok(FindTrailers::Done(len)); 438 } 439 440 let header = temp_buf.get_u8(); 441 442 if header == GRPC_WEB_TRAILERS_BIT { 443 return Ok(FindTrailers::Trailer(len)); 444 } 445 446 if !(header == 0 || header == 1) { 447 return Err(Status::internal(format!( 448 "Invalid header bit {} expected 0 or 1", 449 header 450 ))); 451 } 452 453 let msg_len = temp_buf.get_u32(); 454 455 len += msg_len as usize + 4 + 1; 456 457 // If the msg len of a non-grpc-web trailer frame is larger than 458 // the overall buffer we know within that buffer there are no trailers. 459 if len > buf.len() { 460 return Ok(FindTrailers::IncompleteBuf); 461 } 462 463 temp_buf = &buf[len..]; 464 } 465 } 466 467 #[derive(Debug, PartialEq, Eq)] 468 enum FindTrailers { 469 Trailer(usize), 470 IncompleteBuf, 471 Done(usize), 472 } 473 474 #[cfg(test)] 475 mod tests { 476 use tonic::Code; 477 478 use super::*; 479 480 #[test] 481 fn encoding_constructors() { 482 let cases = &[ 483 (GRPC_WEB, Encoding::None), 484 (GRPC_WEB_PROTO, Encoding::None), 485 (GRPC_WEB_TEXT, Encoding::Base64), 486 (GRPC_WEB_TEXT_PROTO, Encoding::Base64), 487 ("foo", Encoding::None), 488 ]; 489 490 let mut headers = HeaderMap::new(); 491 492 for case in cases { 493 headers.insert(header::CONTENT_TYPE, case.0.parse().unwrap()); 494 headers.insert(header::ACCEPT, case.0.parse().unwrap()); 495 496 assert_eq!(Encoding::from_content_type(&headers), case.1, "{}", case.0); 497 assert_eq!(Encoding::from_accept(&headers), case.1, "{}", case.0); 498 } 499 } 500 501 #[test] 502 fn decode_trailers() { 503 let mut headers = HeaderMap::new(); 504 headers.insert("grpc-status", 0.try_into().unwrap()); 505 headers.insert("grpc-message", "this is a message".try_into().unwrap()); 506 507 let trailers = make_trailers_frame(headers.clone()); 508 509 let buf = Bytes::from(trailers); 510 511 let map = decode_trailers_frame(buf).unwrap().unwrap(); 512 513 assert_eq!(headers, map); 514 } 515 516 #[test] 517 fn find_trailers_non_buffered() { 518 // Byte version of this: 519 // b"\x80\0\0\0\x0fgrpc-status:0\r\n" 520 let buf = [ 521 128, 0, 0, 0, 15, 103, 114, 112, 99, 45, 115, 116, 97, 116, 117, 115, 58, 48, 13, 10, 522 ]; 523 524 let out = find_trailers(&buf[..]).unwrap(); 525 526 assert_eq!(out, FindTrailers::Trailer(0)); 527 } 528 529 #[test] 530 fn find_trailers_buffered() { 531 // Byte version of this: 532 // b"\0\0\0\0L\n$975738af-1a17-4aea-b887-ed0bbced6093\x1a$da609e9b-f470-4cc0-a691-3fd6a005a436\x80\0\0\0\x0fgrpc-status:0\r\n" 533 let buf = [ 534 0, 0, 0, 0, 76, 10, 36, 57, 55, 53, 55, 51, 56, 97, 102, 45, 49, 97, 49, 55, 45, 52, 535 97, 101, 97, 45, 98, 56, 56, 55, 45, 101, 100, 48, 98, 98, 99, 101, 100, 54, 48, 57, 536 51, 26, 36, 100, 97, 54, 48, 57, 101, 57, 98, 45, 102, 52, 55, 48, 45, 52, 99, 99, 48, 537 45, 97, 54, 57, 49, 45, 51, 102, 100, 54, 97, 48, 48, 53, 97, 52, 51, 54, 128, 0, 0, 0, 538 15, 103, 114, 112, 99, 45, 115, 116, 97, 116, 117, 115, 58, 48, 13, 10, 539 ]; 540 541 let out = find_trailers(&buf[..]).unwrap(); 542 543 assert_eq!(out, FindTrailers::Trailer(81)); 544 545 let trailers = decode_trailers_frame(Bytes::copy_from_slice(&buf[81..])) 546 .unwrap() 547 .unwrap(); 548 let status = trailers.get("grpc-status").unwrap(); 549 assert_eq!(status.to_str().unwrap(), "0") 550 } 551 552 #[test] 553 fn find_trailers_buffered_incomplete_message() { 554 let buf = vec![ 555 0, 0, 0, 9, 238, 10, 233, 19, 18, 230, 19, 10, 9, 10, 1, 120, 26, 4, 84, 69, 88, 84, 556 18, 60, 10, 58, 10, 56, 3, 0, 0, 0, 44, 0, 0, 0, 0, 0, 0, 0, 116, 104, 105, 115, 32, 557 118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116, 101, 110, 32, 558 118, 105, 97, 32, 119, 114, 105, 116, 101, 32, 100, 101, 108, 101, 103, 97, 116, 105, 559 111, 110, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, 0, 0, 116, 104, 560 105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116, 561 101, 110, 32, 98, 121, 32, 97, 110, 32, 101, 109, 98, 101, 100, 100, 101, 100, 32, 114, 562 101, 112, 108, 105, 99, 97, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, 563 0, 0, 116, 104, 105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, 564 105, 116, 116, 101, 110, 32, 98, 121, 32, 97, 110, 32, 101, 109, 98, 101, 100, 100, 565 101, 100, 32, 114, 101, 112, 108, 105, 99, 97, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0, 566 46, 0, 0, 0, 0, 0, 0, 0, 116, 104, 105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97, 567 115, 32, 119, 114, 105, 116, 116, 101, 110, 32, 98, 121, 32, 97, 110, 32, 101, 109, 98, 568 101, 100, 100, 101, 100, 32, 114, 101, 112, 108, 105, 99, 97, 33, 18, 62, 10, 60, 10, 569 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, 0, 0, 116, 104, 105, 115, 32, 118, 97, 108, 117, 570 101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116, 101, 110, 32, 98, 121, 32, 97, 110, 571 32, 101, 109, 98, 101, 100, 100, 101, 100, 32, 114, 101, 112, 108, 105, 99, 97, 33, 18, 572 62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, 0, 0, 116, 104, 105, 115, 32, 118, 573 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116, 101, 110, 32, 98, 574 121, 32, 97, 110, 32, 101, 109, 98, 101, 100, 100, 101, 100, 32, 114, 101, 112, 108, 575 105, 99, 97, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, 0, 0, 116, 104, 576 105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, 105, 116, 116, 577 101, 110, 32, 98, 121, 32, 97, 110, 32, 101, 109, 98, 101, 100, 100, 101, 100, 32, 114, 578 101, 112, 108, 105, 99, 97, 33, 18, 62, 10, 60, 10, 58, 3, 0, 0, 0, 46, 0, 0, 0, 0, 0, 579 0, 0, 116, 104, 105, 115, 32, 118, 97, 108, 117, 101, 32, 119, 97, 115, 32, 119, 114, 580 105, 116, 116, 101, 110, 32, 98, 121, 32, 581 ]; 582 583 let out = find_trailers(&buf[..]).unwrap(); 584 585 assert_eq!(out, FindTrailers::IncompleteBuf); 586 } 587 588 #[test] 589 #[ignore] 590 fn find_trailers_buffered_incomplete_buf_bug() { 591 let buf = std::fs::read("tests/incomplete-buf-bug.bin").unwrap(); 592 let out = find_trailers(&buf[..]).unwrap_err(); 593 594 assert_eq!(out.code(), Code::Internal); 595 } 596 } 597