feat: Add metadata to error responses (#348)

tests/integration_tests/tests/status.rs

@@ -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();
+}

tonic/src/status.rs

@@ -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,18 +350,26 @@ 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);
                     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(),
         )
     }
 }