chore(tonic-web): include crate in top-level workspace (#648)

This PR adds the tonic-web crate to tonic's workspace members. Unit and integration tests should now run as part of CI runs.
2021-06-03 11:48:53 -05:00
parent 3e95553c57
commit 8084f4ea26
21 changed files with 28 additions and 425 deletions
@@ -0,0 +1,305 @@
+use std::error::Error;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use bytes::{Buf, BufMut, Bytes, BytesMut};
+use futures_core::{ready, Stream};
+use http::{header, HeaderMap, HeaderValue};
+use http_body::{Body, SizeHint};
+use pin_project::pin_project;
+use tonic::Status;
+
+use self::content_types::*;
+
+pub(crate) mod content_types {
+    use http::{header::CONTENT_TYPE, HeaderMap};
+
+    pub(crate) const GRPC_WEB: &str = "application/grpc-web";
+    pub(crate) const GRPC_WEB_PROTO: &str = "application/grpc-web+proto";
+    pub(crate) const GRPC_WEB_TEXT: &str = "application/grpc-web-text";
+    pub(crate) const GRPC_WEB_TEXT_PROTO: &str = "application/grpc-web-text+proto";
+
+    pub(crate) fn is_grpc_web(headers: &HeaderMap) -> bool {
+        matches!(
+            content_type(headers),
+            Some(GRPC_WEB) | Some(GRPC_WEB_PROTO) | Some(GRPC_WEB_TEXT) | Some(GRPC_WEB_TEXT_PROTO)
+        )
+    }
+
+    fn content_type(headers: &HeaderMap) -> Option<&str> {
+        headers.get(CONTENT_TYPE).and_then(|val| val.to_str().ok())
+    }
+}
+
+const BUFFER_SIZE: usize = 8 * 1024;
+
+const FRAME_HEADER_SIZE: usize = 5;
+
+// 8th (MSB) bit of the 1st gRPC frame byte
+// denotes an uncompressed trailer (as part of the body)
+const GRPC_WEB_TRAILERS_BIT: u8 = 0b10000000;
+
+#[derive(Copy, Clone, PartialEq, Debug)]
+enum Direction {
+    Request,
+    Response,
+}
+
+#[derive(Copy, Clone, PartialEq, Debug)]
+pub(crate) enum Encoding {
+    Base64,
+    None,
+}
+
+#[pin_project]
+pub(crate) struct GrpcWebCall<B> {
+    #[pin]
+    inner: B,
+    buf: BytesMut,
+    direction: Direction,
+    encoding: Encoding,
+    poll_trailers: bool,
+}
+
+impl<B> GrpcWebCall<B> {
+    pub(crate) fn request(inner: B, encoding: Encoding) -> Self {
+        Self::new(inner, Direction::Request, encoding)
+    }
+
+    pub(crate) fn response(inner: B, encoding: Encoding) -> Self {
+        Self::new(inner, Direction::Response, encoding)
+    }
+
+    fn new(inner: B, direction: Direction, encoding: Encoding) -> Self {
+        GrpcWebCall {
+            inner,
+            buf: BytesMut::with_capacity(match (direction, encoding) {
+                (Direction::Response, Encoding::Base64) => BUFFER_SIZE,
+                _ => 0,
+            }),
+            direction,
+            encoding,
+            poll_trailers: true,
+        }
+    }
+
+    // This is to avoid passing a slice of bytes with a length that the base64
+    // decoder would consider invalid.
+    #[inline]
+    fn max_decodable(&self) -> usize {
+        (self.buf.len() / 4) * 4
+    }
+
+    fn decode_chunk(mut self: Pin<&mut Self>) -> Result<Option<Bytes>, Status> {
+        // not enough bytes to decode
+        if self.buf.is_empty() || self.buf.len() < 4 {
+            return Ok(None);
+        }
+
+        // Split `buf` at the largest index that is multiple of 4. Decode the
+        // returned `Bytes`, keeping the rest for the next attempt to decode.
+        let index = self.max_decodable();
+
+        base64::decode(self.as_mut().project().buf.split_to(index))
+            .map(|decoded| Some(Bytes::from(decoded)))
+            .map_err(internal_error)
+    }
+}
+
+impl<B> GrpcWebCall<B>
+where
+    B: Body<Data = Bytes>,
+    B::Error: Error,
+{
+    fn poll_decode(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<B::Data, Status>>> {
+        match self.encoding {
+            Encoding::Base64 => loop {
+                if let Some(bytes) = self.as_mut().decode_chunk()? {
+                    return Poll::Ready(Some(Ok(bytes)));
+                }
+
+                let mut this = self.as_mut().project();
+
+                match ready!(this.inner.as_mut().poll_data(cx)) {
+                    Some(Ok(data)) => this.buf.put(data),
+                    Some(Err(e)) => return Poll::Ready(Some(Err(internal_error(e)))),
+                    None => {
+                        return if this.buf.has_remaining() {
+                            Poll::Ready(Some(Err(internal_error("malformed base64 request"))))
+                        } else {
+                            Poll::Ready(None)
+                        }
+                    }
+                }
+            },
+
+            Encoding::None => match ready!(self.project().inner.poll_data(cx)) {
+                Some(res) => Poll::Ready(Some(res.map_err(internal_error))),
+                None => Poll::Ready(None),
+            },
+        }
+    }
+
+    fn poll_encode(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<B::Data, Status>>> {
+        let mut this = self.as_mut().project();
+
+        if let Some(mut res) = ready!(this.inner.as_mut().poll_data(cx)) {
+            if *this.encoding == Encoding::Base64 {
+                res = res.map(|b| base64::encode(b).into())
+            }
+
+            return Poll::Ready(Some(res.map_err(internal_error)));
+        }
+
+        // this flag is needed because the inner stream never
+        // returns Poll::Ready(None) when polled for trailers
+        if *this.poll_trailers {
+            return match ready!(this.inner.poll_trailers(cx)) {
+                Ok(Some(map)) => {
+                    let mut frame = make_trailers_frame(map);
+
+                    if *this.encoding == Encoding::Base64 {
+                        frame = base64::encode(frame).into_bytes();
+                    }
+
+                    *this.poll_trailers = false;
+                    Poll::Ready(Some(Ok(frame.into())))
+                }
+                Ok(None) => Poll::Ready(None),
+                Err(e) => Poll::Ready(Some(Err(internal_error(e)))),
+            };
+        }
+
+        Poll::Ready(None)
+    }
+}
+
+impl<B> Body for GrpcWebCall<B>
+where
+    B: Body<Data = Bytes>,
+    B::Error: Error,
+{
+    type Data = Bytes;
+    type Error = Status;
+
+    fn poll_data(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Self::Data, Self::Error>>> {
+        match self.direction {
+            Direction::Request => self.poll_decode(cx),
+            Direction::Response => self.poll_encode(cx),
+        }
+    }
+
+    fn poll_trailers(
+        self: Pin<&mut Self>,
+        _: &mut Context<'_>,
+    ) -> Poll<Result<Option<HeaderMap<HeaderValue>>, Self::Error>> {
+        Poll::Ready(Ok(None))
+    }
+
+    fn is_end_stream(&self) -> bool {
+        self.inner.is_end_stream()
+    }
+
+    fn size_hint(&self) -> SizeHint {
+        self.inner.size_hint()
+    }
+}
+
+impl<B> Stream for GrpcWebCall<B>
+where
+    B: Body<Data = Bytes>,
+    B::Error: Error,
+{
+    type Item = Result<Bytes, Status>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        Body::poll_data(self, cx)
+    }
+}
+
+impl Encoding {
+    pub(crate) fn from_content_type(headers: &HeaderMap) -> Encoding {
+        Self::from_header(headers.get(header::CONTENT_TYPE))
+    }
+
+    pub(crate) fn from_accept(headers: &HeaderMap) -> Encoding {
+        Self::from_header(headers.get(header::ACCEPT))
+    }
+
+    pub(crate) fn to_content_type(&self) -> &'static str {
+        match self {
+            Encoding::Base64 => GRPC_WEB_TEXT_PROTO,
+            Encoding::None => GRPC_WEB_PROTO,
+        }
+    }
+
+    fn from_header(value: Option<&HeaderValue>) -> Encoding {
+        match value.and_then(|val| val.to_str().ok()) {
+            Some(GRPC_WEB_TEXT_PROTO) | Some(GRPC_WEB_TEXT) => Encoding::Base64,
+            _ => Encoding::None,
+        }
+    }
+}
+
+fn internal_error(e: impl std::fmt::Display) -> Status {
+    Status::internal(format!("tonic-web: {}", e))
+}
+
+// Key-value pairs encoded as a HTTP/1 headers block (without the terminating newline)
+fn encode_trailers(trailers: HeaderMap) -> Vec<u8> {
+    trailers.iter().fold(Vec::new(), |mut acc, (key, value)| {
+        acc.put_slice(key.as_ref());
+        acc.push(b':');
+        acc.put_slice(value.as_bytes());
+        acc.put_slice(b"\r\n");
+        acc
+    })
+}
+
+fn make_trailers_frame(trailers: HeaderMap) -> Vec<u8> {
+    let trailers = encode_trailers(trailers);
+    let len = trailers.len();
+    assert!(len <= u32::MAX as usize);
+
+    let mut frame = Vec::with_capacity(len + FRAME_HEADER_SIZE);
+    frame.push(GRPC_WEB_TRAILERS_BIT);
+    frame.put_u32(len as u32);
+    frame.extend(trailers);
+
+    frame
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn encoding_constructors() {
+        let cases = &[
+            (GRPC_WEB, Encoding::None),
+            (GRPC_WEB_PROTO, Encoding::None),
+            (GRPC_WEB_TEXT, Encoding::Base64),
+            (GRPC_WEB_TEXT_PROTO, Encoding::Base64),
+            ("foo", Encoding::None),
+        ];
+
+        let mut headers = HeaderMap::new();
+
+        for case in cases {
+            headers.insert(header::CONTENT_TYPE, case.0.parse().unwrap());
+            headers.insert(header::ACCEPT, case.0.parse().unwrap());
+
+            assert_eq!(Encoding::from_content_type(&headers), case.1, "{}", case.0);
+            assert_eq!(Encoding::from_accept(&headers), case.1, "{}", case.0);
+        }
+    }
+}