diff --git a/tests/integration_tests/tests/status.rs b/tests/integration_tests/tests/status.rs index cf750ac8d..d3e69f13e 100644 --- a/tests/integration_tests/tests/status.rs +++ b/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) -> Result, 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(); +} diff --git a/tonic/src/status.rs b/tonic/src/status.rs index 9ee6c73f9..2cc4515bc 100644 --- a/tonic/src/status.rs +++ b/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,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 { - 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, 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, 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, + 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(), ) } }