feat: Add metadata to error responses (#348)

2020-05-10 18:46:45 +03:00
parent 8cd6f05064
commit 372da52e96
2 changed files with 125 additions and 5 deletions
@@ -3,6 +3,7 @@ use futures_util::FutureExt;
 use integration_tests::pb::{test_client, test_server, Input, Output};
 use std::time::Duration;
 use tokio::sync::oneshot;
+use tonic::metadata::{MetadataMap, MetadataValue};
 use tonic::{transport::Server, Code, Request, Response, Status};

 #[tokio::test]
@@ -50,3 +51,69 @@ async fn status_with_details() {

    jh.await.unwrap();
 }
+
+#[tokio::test]
+async fn status_with_metadata() {
+    const MESSAGE: &str = "Internal error, see metadata for details";
+
+    const ASCII_NAME: &str = "x-host-ip";
+    const ASCII_VALUE: &str = "127.0.0.1";
+
+    const BINARY_NAME: &str = "x-host-name-bin";
+    const BINARY_VALUE: &[u8] = b"localhost";
+
+    struct Svc;
+
+    #[tonic::async_trait]
+    impl test_server::Test for Svc {
+        async fn unary_call(&self, _: Request<Input>) -> Result<Response<Output>, Status> {
+            let mut metadata = MetadataMap::new();
+            metadata.insert(ASCII_NAME, ASCII_VALUE.parse().unwrap());
+            metadata.insert_bin(BINARY_NAME, MetadataValue::from_bytes(BINARY_VALUE));
+
+            Err(Status::with_metadata(Code::Internal, MESSAGE, metadata))
+        }
+    }
+
+    let svc = test_server::TestServer::new(Svc);
+
+    let (tx, rx) = oneshot::channel::<()>();
+
+    let jh = tokio::spawn(async move {
+        Server::builder()
+            .add_service(svc)
+            .serve_with_shutdown("127.0.0.1:1338".parse().unwrap(), rx.map(drop))
+            .await
+            .unwrap();
+    });
+
+    tokio::time::delay_for(Duration::from_millis(100)).await;
+
+    let mut channel = test_client::TestClient::connect("http://127.0.0.1:1338")
+        .await
+        .unwrap();
+
+    let err = channel
+        .unary_call(Request::new(Input {}))
+        .await
+        .unwrap_err();
+
+    assert_eq!(err.code(), Code::Internal);
+    assert_eq!(err.message(), MESSAGE);
+
+    let metadata = err.metadata();
+
+    assert_eq!(
+        metadata.get(ASCII_NAME).unwrap().to_str().unwrap(),
+        ASCII_VALUE
+    );
+
+    assert_eq!(
+        metadata.get_bin(BINARY_NAME).unwrap().to_bytes().unwrap(),
+        BINARY_VALUE
+    );
+
+    tx.send(()).unwrap();
+
+    jh.await.unwrap();
+}
@@ -1,3 +1,4 @@
+use crate::metadata::MetadataMap;
 use bytes::Bytes;
 use http::header::{HeaderMap, HeaderValue};
 use percent_encoding::{percent_decode, percent_encode, AsciiSet, CONTROLS};
@@ -39,6 +40,10 @@ pub struct Status {
    message: String,
    /// Binary opaque details, found in the `grpc-status-details-bin` header.
    details: Bytes,
+    /// Custom metadata, found in the user-defined headers.
+    /// If the metadata contains any headers with names reserved either by the gRPC spec
+    /// or by `Status` fields above, they will be ignored.
+    metadata: MetadataMap,
 }

 /// gRPC status codes used by [`Status`].
@@ -113,6 +118,7 @@ impl Status {
            code,
            message: message.into(),
            details: Bytes::new(),
+            metadata: MetadataMap::new(),
        }
    }

@@ -266,6 +272,7 @@ impl Status {
                    code: status.code,
                    message: status.message.clone(),
                    details: status.details.clone(),
+                    metadata: status.metadata.clone(),
                });
            }

@@ -343,11 +350,18 @@ impl Status {
                })
                .map(Bytes::from)
                .unwrap_or_else(Bytes::new);
+
+            let mut other_headers = header_map.clone();
+            other_headers.remove(GRPC_STATUS_HEADER_CODE);
+            other_headers.remove(GRPC_STATUS_MESSAGE_HEADER);
+            other_headers.remove(GRPC_STATUS_DETAILS_HEADER);
+
            match error_message {
                Ok(message) => Status {
                    code,
                    message,
                    details,
+                    metadata: MetadataMap::from_headers(other_headers),
                },
                Err(err) => {
                    warn!("Error deserializing status message header: {}", err);
@@ -355,6 +369,7 @@ impl Status {
                        code: Code::Unknown,
                        message: format!("Error deserializing status message header: {}", err),
                        details,
+                        metadata: MetadataMap::from_headers(other_headers),
                    }
                }
            }
@@ -376,13 +391,25 @@ impl Status {
        &self.details
    }

+    /// Get a reference to the custom metadata.
+    pub fn metadata(&self) -> &MetadataMap {
+        &self.metadata
+    }
+
+    /// Get a mutable reference to the custom metadata.
+    pub fn metadata_mut(&mut self) -> &mut MetadataMap {
+        &mut self.metadata
+    }
+
    pub(crate) fn to_header_map(&self) -> Result<HeaderMap, Self> {
-        let mut header_map = HeaderMap::with_capacity(3);
+        let mut header_map = HeaderMap::with_capacity(3 + self.metadata.len());
        self.add_header(&mut header_map)?;
        Ok(header_map)
    }

    pub(crate) fn add_header(&self, header_map: &mut HeaderMap) -> Result<(), Self> {
+        header_map.extend(self.metadata.clone().into_sanitized_headers());
+
        header_map.insert(GRPC_STATUS_HEADER_CODE, self.code.to_header_value());

        if !self.message.is_empty() {
@@ -409,12 +436,32 @@ impl Status {

    /// Create a new `Status` with the associated code, message, and binary details field.
    pub fn with_details(code: Code, message: impl Into<String>, details: Bytes) -> Status {
-        let details = base64::encode_config(&details[..], base64::STANDARD_NO_PAD);
+        Self::with_details_and_metadata(code, message, details, MetadataMap::new())
+    }
+
+    /// Create a new `Status` with the associated code, message, and custom metadata
+    pub fn with_metadata(code: Code, message: impl Into<String>, metadata: MetadataMap) -> Status {
+        Self::with_details_and_metadata(code, message, Bytes::new(), metadata)
+    }
+
+    /// Create a new `Status` with the associated code, message, binary details field and custom metadata
+    pub fn with_details_and_metadata(
+        code: Code,
+        message: impl Into<String>,
+        details: Bytes,
+        metadata: MetadataMap,
+    ) -> Status {
+        let details = if details.is_empty() {
+            details
+        } else {
+            base64::encode_config(&details[..], base64::STANDARD_NO_PAD).into()
+        };

        Status {
            code,
            message: message.into(),
-            details: details.into(),
+            details: details,
+            metadata: metadata,
        }
    }
 }
@@ -434,6 +481,10 @@ impl fmt::Debug for Status {
            builder.field("details", &self.details);
        }

+        if !self.metadata.is_empty() {
+            builder.field("metadata", &self.metadata);
+        }
+
        builder.finish()
    }
 }
@@ -470,9 +521,11 @@ impl fmt::Display for Status {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
-            "grpc-status: {:?}, grpc-message: {:?}",
+            "status: {:?}, message: {:?}, details: {:?}, metadata: {:?}",
            self.code(),
-            self.message()
+            self.message(),
+            self.details(),
+            self.metadata(),
        )
    }
 }