Extensions in Status

tonic/src/response.rs

@@ -61,6 +61,16 @@ impl<T> Response<T> {
         (self.metadata, self.message, self.extensions)
     }
 
+    /// Create a new gRPC response from massage and metadata.
+    pub fn with_metadata(message: T, metadata: MetadataMap) -> Self {
+        Self::from_parts(metadata, message, Extensions::new())
+    }
+
+    /// Create a new gRPC response from message and extensions.
+    pub fn with_extensions(message: T, extensions: Extensions) -> Self {
+        Self::from_parts(MetadataMap::new(), message, extensions)
+    }
+
     /// Create a new gRPC response from metadata, message and extensions.
     pub fn from_parts(metadata: MetadataMap, message: T, extensions: Extensions) -> Self {
         Self {

tonic/src/status.rs

@@ -2,6 +2,7 @@ use crate::metadata::MetadataMap;
 use crate::metadata::GRPC_CONTENT_TYPE;
 use base64::Engine as _;
 use bytes::Bytes;
+use http::Extensions;
 use http::{
     header::{HeaderMap, HeaderValue},
     HeaderName,
@@ -50,6 +51,8 @@ struct StatusInner {
     /// 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,
+    /// Data for passing information from the Service to middleware
+    extensions: Extensions,
     /// Optional underlying error.
     source: Option<Arc<dyn Error + Send + Sync + 'static>>,
 }
@@ -175,6 +178,7 @@ impl Status {
             message: message.into(),
             details: Bytes::new(),
             metadata: MetadataMap::new(),
+            extensions: Extensions::new(),
             source: None,
         }
         .into_status()
@@ -502,6 +506,7 @@ impl Status {
                 message,
                 details,
                 metadata: MetadataMap::from_headers(other_headers),
+                extensions: Extensions::new(),
                 source: None,
             }
             .into_status(),
@@ -533,6 +538,16 @@ impl Status {
         &mut self.0.metadata
     }
 
+    /// Get a reference to the custom extensions.
+    pub fn extensions(&self) -> &Extensions {
+        &self.0.extensions
+    }
+
+    /// Get a mutable reference to the custom extensions.
+    pub fn extensions_mut(&mut self) -> &mut Extensions {
+        &mut self.0.extensions
+    }
+
     pub(crate) fn to_header_map(&self) -> Result<HeaderMap, Self> {
         let mut header_map = HeaderMap::with_capacity(3 + self.0.metadata.len());
         self.add_header(&mut header_map)?;
@@ -590,6 +605,24 @@ impl Status {
             message: message.into(),
             details,
             metadata,
+            extensions: Extensions::new(),
+            source: None,
+        }
+        .into_status()
+    }
+
+    /// Create a new `Status` with the associated code, message, and custom extensions
+    pub fn with_extensions(
+        code: Code,
+        message: impl Into<String>,
+        extensions: Extensions,
+    ) -> Status {
+        StatusInner {
+            code,
+            message: message.into(),
+            details: Bytes::new(),
+            metadata: MetadataMap::new(),
+            extensions,
             source: None,
         }
         .into_status()
@@ -608,7 +641,7 @@ impl Status {
             .headers_mut()
             .insert(http::header::CONTENT_TYPE, GRPC_CONTENT_TYPE);
         self.add_header(response.headers_mut()).unwrap();
-        response.extensions_mut().insert(self);
+        response.extensions_mut().extend(self.0.extensions);
         response
     }
 
@@ -631,6 +664,7 @@ fn find_status_in_source_chain(err: &(dyn Error + 'static)) -> Option<Status> {
                     message: status.0.message.clone(),
                     details: status.0.details.clone(),
                     metadata: status.0.metadata.clone(),
+                    extensions: status.0.extensions.clone(),
                     // Since `Status` is not `Clone`, any `source` on the original Status
                     // cannot be cloned so must remain with the original `Status`.
                     source: None,