Headers passthrough

crates/ingress-dispatcher/src/dispatcher.rs

@@ -92,6 +92,7 @@ impl DispatchIngressRequest for IngressDispatcher {
             span_context,
             request_mode,
             idempotency,
+            headers,
         } = ingress_request;
 
         let invocation_id: InvocationId = fid.clone().into();
@@ -134,6 +135,7 @@ impl DispatchIngressRequest for IngressDispatcher {
                         source: invocation::Source::Ingress,
                         response_sink,
                         span_context,
+                        headers,
                     },
                     MapResponseAction::IdempotentInvokerResponse,
                 )
@@ -146,6 +148,7 @@ impl DispatchIngressRequest for IngressDispatcher {
                         source: invocation::Source::Ingress,
                         response_sink,
                         span_context,
+                        headers,
                     },
                     MapResponseAction::None,
                 )
@@ -314,6 +317,7 @@ mod tests {
                 argument.clone(),
                 SpanRelation::None,
                 IdempotencyMode::key(idempotency_key.clone(), None),
+                vec![],
             );
             dispatcher.dispatch_ingress_request(invocation).await?;
 

crates/ingress-dispatcher/src/lib.rs

@@ -55,6 +55,7 @@ pub struct IngressRequest {
     span_context: ServiceInvocationSpanContext,
     request_mode: IngressRequestMode,
     idempotency: IdempotencyMode,
+    headers: Vec<restate_types::invocation::Header>,
 }
 
 #[derive(Debug, Clone)]
@@ -104,6 +105,7 @@ impl IngressRequest {
         argument: impl Into<Bytes>,
         related_span: SpanRelation,
         idempotency: IdempotencyMode,
+        headers: Vec<restate_types::invocation::Header>,
     ) -> (Self, IngressResponseReceiver) {
         let span_context = ServiceInvocationSpanContext::start(&fid, related_span);
         let (result_tx, result_rx) = oneshot::channel();
@@ -116,6 +118,7 @@ impl IngressRequest {
                 request_mode: IngressRequestMode::RequestResponse(result_tx),
                 span_context,
                 idempotency,
+                headers,
             },
             result_rx,
         )
@@ -127,6 +130,7 @@ impl IngressRequest {
         argument: impl Into<Bytes>,
         related_span: SpanRelation,
         ingress_deduplication_id: Option<IngressDeduplicationId>,
+        headers: Vec<restate_types::invocation::Header>,
     ) -> Self {
         let span_context = ServiceInvocationSpanContext::start(&fid, related_span);
         IngressRequest {
@@ -139,6 +143,7 @@ impl IngressRequest {
                 Some(dedup_id) => IngressRequestMode::DedupFireAndForget(dedup_id),
             },
             idempotency: IdempotencyMode::None,
+            headers,
         }
     }
 
@@ -147,6 +152,7 @@ impl IngressRequest {
         event: Event,
         related_span: SpanRelation,
         deduplication: Option<(D, MessageIndex)>,
+        headers: Vec<restate_types::invocation::Header>,
     ) -> Result<Self, anyhow::Error> {
         // Check if we need to proxy or not
         let (proxying_key, request_mode) = if let Some((dedup_id, dedup_index)) = deduplication {
@@ -213,6 +219,7 @@ impl IngressRequest {
                 span_context,
                 request_mode,
                 idempotency: IdempotencyMode::None,
+                headers,
             }
         } else {
             IngressRequest {
@@ -222,6 +229,7 @@ impl IngressRequest {
                 span_context,
                 request_mode,
                 idempotency: IdempotencyMode::None,
+                headers,
             }
         })
     }
@@ -290,6 +298,7 @@ pub mod mocks {
             ServiceInvocationSpanContext,
             IdempotencyMode,
             IngressResponseSender,
+            Vec<restate_types::invocation::Header>,
         ) {
             let_assert!(
                 IngressRequest {
@@ -299,6 +308,7 @@ pub mod mocks {
                     span_context,
                     request_mode: IngressRequestMode::RequestResponse(ingress_response_sender),
                     idempotency,
+                    headers,
                     ..
                 } = self
             );
@@ -309,6 +319,7 @@ pub mod mocks {
                 span_context,
                 idempotency,
                 ingress_response_sender,
+                headers,
             )
         }
 

crates/ingress-http/src/handler/awakeables.rs

@@ -96,6 +96,7 @@ where
                 payload,
                 SpanRelation::Linked(ingress_span_context),
                 IdempotencyMode::None,
+                vec![],
             );
             if let Err(e) = dispatcher.dispatch_ingress_request(invocation).await {
                 warn!(

crates/ingress-http/src/handler/component_handler.rs

@@ -21,7 +21,7 @@ use metrics::{counter, histogram};
 use restate_ingress_dispatcher::{DispatchIngressRequest, IdempotencyMode, IngressRequest};
 use restate_schema_api::component::ComponentMetadataResolver;
 use restate_types::identifiers::{FullInvocationId, InvocationId, ServiceId};
-use restate_types::invocation::SpanRelation;
+use restate_types::invocation::{Header, SpanRelation};
 use serde::Serialize;
 use std::num::ParseIntError;
 use std::time::{Duration, Instant};
@@ -86,20 +86,24 @@ where
         let handle_fut = async move {
             info!("Processing ingress request");
 
+            let (parts, body) = req.into_parts();
+
             // Check HTTP Method
-            if req.method() != Method::GET && req.method() != Method::POST {
+            if parts.method != Method::GET && parts.method != Method::POST {
                 return Err(HandlerError::MethodNotAllowed);
             }
 
             // TODO validate content-type
             //  https://github.com/restatedev/restate/issues/1230
 
             // Check if Idempotency-Key is available
-            let idempotency_mode = parse_idempotency_key_and_retention_period(req.headers())?;
+            let idempotency_mode = parse_idempotency_key_and_retention_period(&parts.headers)?;
+
+            // Get headers
+            let headers = parse_headers(parts.headers)?;
 
             // Collect body
-            let collected_request_bytes = req
-                .into_body()
+            let collected_request_bytes = body
                 .collect()
                 .await
                 .map_err(|e| HandlerError::Body(e.into()))?
@@ -116,6 +120,7 @@ where
                         idempotency_mode,
                         collected_request_bytes,
                         span_relation,
+                        headers,
                         self.dispatcher,
                     )
                     .await
@@ -127,6 +132,7 @@ where
                         idempotency_mode,
                         collected_request_bytes,
                         span_relation,
+                        headers,
                         self.dispatcher,
                     )
                     .await
@@ -164,6 +170,7 @@ where
         idempotency_mode: IdempotencyMode,
         collected_request_bytes: Bytes,
         span_relation: SpanRelation,
+        headers: Vec<Header>,
         dispatcher: Dispatcher,
     ) -> Result<Response<Full<Bytes>>, HandlerError> {
         let invocation_id: InvocationId = fid.clone().into();
@@ -173,6 +180,7 @@ where
             collected_request_bytes,
             span_relation,
             idempotency_mode,
+            headers,
         );
         if let Err(e) = dispatcher.dispatch_ingress_request(invocation).await {
             warn!(
@@ -222,6 +230,7 @@ where
         idempotency_mode: IdempotencyMode,
         collected_request_bytes: Bytes,
         span_relation: SpanRelation,
+        headers: Vec<Header>,
         dispatcher: Dispatcher,
     ) -> Result<Response<Full<Bytes>>, HandlerError> {
         if !matches!(idempotency_mode, IdempotencyMode::None) {
@@ -238,6 +247,7 @@ where
             collected_request_bytes,
             span_relation,
             None,
+            headers,
         );
         if let Err(e) = dispatcher.dispatch_ingress_request(invocation).await {
             warn!(
@@ -261,6 +271,27 @@ where
     }
 }
 
+fn parse_headers(headers: HeaderMap) -> Result<Vec<Header>, HandlerError> {
+    headers
+        .into_iter()
+        .filter_map(|(k, v)| k.map(|k| (k, v)))
+        // Filter out Connection, Host and idempotency headers
+        .filter(|(k, _)| {
+            k != header::CONNECTION
+                && k != header::HOST
+                && k != IDEMPOTENCY_KEY
+                && k != IDEMPOTENCY_EXPIRES
+                && k != IDEMPOTENCY_RETENTION_PERIOD
+        })
+        .map(|(k, v)| {
+            let value = v
+                .to_str()
+                .map_err(|e| HandlerError::BadHeader(k.clone(), e))?;
+            Ok(Header::new(k.as_str(), value))
+        })
+        .collect()
+}
+
 fn parse_idempotency_key_and_retention_period(
     headers: &HeaderMap,
 ) -> Result<IdempotencyMode, HandlerError> {

crates/ingress-http/src/handler/error.rs

@@ -11,7 +11,7 @@
 use super::APPLICATION_JSON;
 
 use bytes::Bytes;
-use http::{Response, StatusCode};
+use http::{header, Response, StatusCode};
 use restate_types::errors::{InvocationError, UserErrorCode};
 use serde::Serialize;
 use std::string;
@@ -30,6 +30,8 @@ pub(crate) enum HandlerError {
     BadAwakeablesPath,
     #[error("not implemented")]
     NotImplemented,
+    #[error("bad header {0}: {1:?}")]
+    BadHeader(header::HeaderName, #[source] header::ToStrError),
     #[error("bad path, cannot decode key: {0:?}")]
     UrlDecodingError(string::FromUtf8Error),
     #[error("the invoked component is not public")]
@@ -82,6 +84,7 @@ impl HandlerError {
             HandlerError::Invocation(e) => {
                 invocation_status_code_to_http_status_code(e.code().into())
             }
+            HandlerError::BadHeader(_, _) => StatusCode::BAD_REQUEST,
         };
 
         let error_response = match self {

crates/ingress-http/src/handler/tests.rs

@@ -22,6 +22,7 @@ use restate_core::TestCoreEnv;
 use restate_ingress_dispatcher::mocks::MockDispatcher;
 use restate_ingress_dispatcher::IdempotencyMode;
 use restate_ingress_dispatcher::IngressRequest;
+use restate_types::invocation::Header;
 use tokio::sync::mpsc;
 use tower::ServiceExt;
 use tracing_test::traced_test;
@@ -36,16 +37,20 @@ async fn call_service() {
     let req = hyper::Request::builder()
         .uri("http://localhost/greeter.Greeter/greet")
         .method(Method::POST)
+        .header("Connection", "Close")
+        .header("my-header", "my-value")
         .body(Full::new(Bytes::from(
             serde_json::to_vec(&greeting_req).unwrap(),
         )))
         .unwrap();
 
     let response = handle(req, |ingress_req| {
         // Get the function invocation and assert on it
-        let (fid, method_name, argument, _, _, response_tx) = ingress_req.expect_invocation();
+        let (fid, method_name, argument, _, _, response_tx, headers) =
+            ingress_req.expect_invocation();
         restate_test_util::assert_eq!(fid.service_id.service_name, "greeter.Greeter");
         restate_test_util::assert_eq!(method_name, "greet");
+        restate_test_util::assert_eq!(headers, vec![Header::new("my-header", "my-value")]);
 
         let greeting_req: GreetingRequest = serde_json::from_slice(&argument).unwrap();
         restate_test_util::assert_eq!(&greeting_req.person, "Francesco");
@@ -81,7 +86,7 @@ async fn call_service_with_get() {
 
     let response = handle(req, |ingress_req| {
         // Get the function invocation and assert on it
-        let (fid, method_name, argument, _, _, response_tx) = ingress_req.expect_invocation();
+        let (fid, method_name, argument, _, _, response_tx, _) = ingress_req.expect_invocation();
         restate_test_util::assert_eq!(fid.service_id.service_name, "greeter.Greeter");
         restate_test_util::assert_eq!(method_name, "greet");
 
@@ -124,7 +129,7 @@ async fn call_virtual_object() {
 
     let response = handle(req, |ingress_req| {
         // Get the function invocation and assert on it
-        let (fid, method_name, argument, _, _, response_tx) = ingress_req.expect_invocation();
+        let (fid, method_name, argument, _, _, response_tx, _) = ingress_req.expect_invocation();
         restate_test_util::assert_eq!(fid.service_id.service_name, "greeter.GreeterObject");
         restate_test_util::assert_eq!(fid.service_id.key, &"my-key");
         restate_test_util::assert_eq!(method_name, "greet");
@@ -235,7 +240,7 @@ async fn idempotency_key_parsing() {
 
     let response = handle(req, |ingress_req| {
         // Get the function invocation and assert on it
-        let (fid, method_name, argument, _, idempotency_mode, response_tx) =
+        let (fid, method_name, argument, _, idempotency_mode, response_tx, _) =
             ingress_req.expect_invocation();
         restate_test_util::assert_eq!(fid.service_id.service_name, "greeter.Greeter");
         restate_test_util::assert_eq!(method_name, "greet");

crates/ingress-http/src/server.rs

@@ -240,7 +240,7 @@ mod tests {
         let (address, input, handle) = bootstrap_test().await;
         let process_fut = tokio::task::spawn(async move {
             // Get the function invocation and assert on it
-            let (fid, method_name, argument, _, _, response_tx) =
+            let (fid, method_name, argument, _, _, response_tx, _) =
                 input.await.unwrap().unwrap().expect_invocation();
             assert_eq!(fid.service_id.service_name, "greeter.Greeter");
             assert_eq!(method_name, "greet");

crates/ingress-kafka/src/consumer_task.rs

@@ -135,6 +135,7 @@ impl MessageSender {
             event,
             SpanRelation::Parent(ingress_span_context),
             Some(Self::generate_deduplication_id(consumer_group_id, msg)),
+            vec![],
         )
         .map_err(|cause| Error::Event {
             topic: msg.topic().to_string(),

crates/service-protocol/src/codec.rs

@@ -12,6 +12,7 @@ use super::pb::protocol;
 
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use prost::Message;
+use restate_types::invocation::Header;
 use restate_types::journal::enriched::{EnrichedEntryHeader, EnrichedRawEntry};
 use restate_types::journal::raw::*;
 use restate_types::journal::{CompletionResult, Entry, EntryType};
@@ -39,11 +40,17 @@ macro_rules! match_decode {
 pub struct ProtobufRawEntryCodec;
 
 impl RawEntryCodec for ProtobufRawEntryCodec {
-    fn serialize_as_input_entry(value: Bytes) -> EnrichedRawEntry {
+    fn serialize_as_input_entry(headers: Vec<Header>, value: Bytes) -> EnrichedRawEntry {
         RawEntry::new(
             EnrichedEntryHeader::Input {},
             protocol::InputEntryMessage {
-                headers: vec![],
+                headers: headers
+                    .into_iter()
+                    .map(|h| protocol::Header {
+                        key: h.name.to_string(),
+                        value: h.value.to_string(),
+                    })
+                    .collect(),
                 value,
             }
             .encode_to_vec()