fix: experimental headers (#1524)

Author: tusharmath

benches/impl_path_string_for_evaluation_context.rs

@@ -253,7 +253,7 @@ fn request_context() -> RequestContext {
 }
 
 fn bench_main(c: &mut Criterion) {
-    let mut req_ctx = request_context().req_headers(TEST_HEADERS.clone());
+    let mut req_ctx = request_context().request_headers(TEST_HEADERS.clone());
 
     req_ctx.server.vars = TEST_VARS.clone();
     let eval_ctx = EvaluationContext::new(&req_ctx, &MockGraphqlContext);

src/blueprint/server.rs

@@ -1,5 +1,6 @@
-use std::collections::{BTreeMap, BTreeSet};
+use std::collections::{BTreeMap, BTreeSet, HashSet};
 use std::net::{AddrParseError, IpAddr};
+use std::str::FromStr;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -33,7 +34,7 @@ pub struct Server {
     pub pipeline_flush: bool,
     pub script: Option<Script>,
     pub cors: Option<Cors>,
-    pub experimental_headers: BTreeSet<String>,
+    pub experimental_headers: HashSet<HeaderName>,
 }
 
 /// Mimic of mini_v8::Script that's wasm compatible
@@ -75,7 +76,7 @@ impl Server {
         self.enable_query_validation
     }
 
-    pub fn get_experimental_headers(&self) -> BTreeSet<String> {
+    pub fn get_experimental_headers(&self) -> HashSet<HeaderName> {
         self.experimental_headers.clone()
     }
 }
@@ -209,7 +210,7 @@ fn handle_response_headers(resp_headers: Vec<(String, String)>) -> Valid<HeaderM
     .trace("schema")
 }
 
-fn handle_experimental_headers(headers: BTreeSet<String>) -> Valid<BTreeSet<String>, String> {
+fn handle_experimental_headers(headers: BTreeSet<String>) -> Valid<HashSet<HeaderName>, String> {
     Valid::from_iter(headers.iter(), |h| {
         if !h.to_lowercase().starts_with("x-") {
             Valid::fail(
@@ -220,10 +221,10 @@ fn handle_experimental_headers(headers: BTreeSet<String>) -> Valid<BTreeSet<Stri
                 .to_string(),
             )
         } else {
-            Valid::succeed(h.clone())
+            Valid::from(HeaderName::from_str(h).map_err(|e| ValidationError::new(e.to_string())))
         }
     })
-    .map_to(headers)
+    .map(HashSet::from_iter)
     .trace("experimental")
     .trace("headers")
     .trace("@server")

src/http/request_context.rs

@@ -20,8 +20,8 @@ use crate::runtime::TargetRuntime;
 pub struct RequestContext {
     pub server: Server,
     pub upstream: Upstream,
-    pub req_headers: HeaderMap,
-    pub experimental_headers: HeaderMap,
+    pub request_headers: HeaderMap,
+    pub x_response_headers: Arc<Mutex<HeaderMap>>,
     pub cookie_headers: Option<Arc<Mutex<HeaderMap>>>,
     pub http_data_loaders: Arc<Vec<DataLoader<DataLoaderRequest, HttpDataLoader>>>,
     pub gql_data_loaders: Arc<Vec<DataLoader<DataLoaderRequest, GraphqlDataLoader>>>,
@@ -37,8 +37,8 @@ impl RequestContext {
         RequestContext {
             server: Default::default(),
             upstream: Default::default(),
-            req_headers: HeaderMap::new(),
-            experimental_headers: HeaderMap::new(),
+            request_headers: HeaderMap::new(),
+            x_response_headers: Arc::new(Mutex::new(HeaderMap::new())),
             cookie_headers: None,
             http_data_loaders: Arc::new(vec![]),
             gql_data_loaders: Arc::new(vec![]),
@@ -145,6 +145,33 @@ impl RequestContext {
     pub fn is_batching_enabled(&self) -> bool {
         self.upstream.is_batching_enabled()
     }
+
+    /// Checks if experimental headers is enabled
+    pub fn has_experimental_headers(&self) -> bool {
+        !self.server.experimental_headers.is_empty()
+    }
+
+    /// Inserts the experimental headers into the x_response_headers map
+    pub fn add_x_headers(&self, headers: &HeaderMap) {
+        if self.has_experimental_headers() {
+            let mut x_response_headers = self.x_response_headers.lock().unwrap();
+            for name in &self.server.experimental_headers {
+                if let Some(value) = headers.get(name) {
+                    x_response_headers.insert(name, value.clone());
+                }
+            }
+        }
+    }
+
+    /// Modifies existing headers to include the experimental headers
+    pub fn extend_x_headers(&self, headers: &mut HeaderMap) {
+        if self.has_experimental_headers() {
+            let x_response_headers = &self.x_response_headers.lock().unwrap();
+            for (header, value) in x_response_headers.iter() {
+                headers.insert(header, value.clone());
+            }
+        }
+    }
 }
 
 impl From<&AppContext> for RequestContext {
@@ -157,8 +184,8 @@ impl From<&AppContext> for RequestContext {
         Self {
             server: app_ctx.blueprint.server.clone(),
             upstream: app_ctx.blueprint.upstream.clone(),
-            req_headers: HeaderMap::new(),
-            experimental_headers: HeaderMap::new(),
+            request_headers: HeaderMap::new(),
+            x_response_headers: Arc::new(Mutex::new(HeaderMap::new())),
             cookie_headers,
             http_data_loaders: app_ctx.http_data_loaders.clone(),
             gql_data_loaders: app_ctx.gql_data_loaders.clone(),

src/http/request_handler.rs

@@ -1,5 +1,6 @@
 use std::borrow::Cow;
 use std::collections::BTreeSet;
+use std::ops::Deref;
 use std::sync::Arc;
 
 use anyhow::Result;
@@ -75,11 +76,8 @@ fn create_request_context(req: &Request<Body>, app_ctx: &AppContext) -> RequestC
     let allowed = upstream.allowed_headers;
     let req_headers = create_allowed_headers(req.headers(), &allowed);
 
-    let allowed = app_ctx.blueprint.server.get_experimental_headers();
-    let experimental_headers = create_allowed_headers(req.headers(), &allowed);
-    RequestContext::from(app_ctx)
-        .req_headers(req_headers)
-        .experimental_headers(experimental_headers)
+    let _allowed = app_ctx.blueprint.server.get_experimental_headers();
+    RequestContext::from(app_ctx).request_headers(req_headers)
 }
 
 fn update_cache_control_header(
@@ -95,30 +93,25 @@ fn update_cache_control_header(
     response
 }
 
-fn update_experimental_headers(
-    response: &mut hyper::Response<hyper::Body>,
-    app_ctx: &AppContext,
-    req_ctx: Arc<RequestContext>,
-) {
-    if !app_ctx.blueprint.server.experimental_headers.is_empty() {
-        response
-            .headers_mut()
-            .extend(req_ctx.experimental_headers.clone());
-    }
-}
-
 pub fn update_response_headers(
     resp: &mut hyper::Response<hyper::Body>,
-    cookie_headers: Option<HeaderMap>,
+    req_ctx: &RequestContext,
     app_ctx: &AppContext,
 ) {
     if !app_ctx.blueprint.server.response_headers.is_empty() {
+        // Add static response headers
         resp.headers_mut()
             .extend(app_ctx.blueprint.server.response_headers.clone());
     }
-    if let Some(cookie_headers) = cookie_headers {
-        resp.headers_mut().extend(cookie_headers);
+
+    // Insert Cookie Headers
+    if let Some(ref cookie_headers) = req_ctx.cookie_headers {
+        let cookie_headers = cookie_headers.lock().unwrap();
+        resp.headers_mut().extend(cookie_headers.deref().clone());
     }
+
+    // Insert Experimental Headers
+    req_ctx.extend_x_headers(resp.headers_mut());
 }
 
 #[tracing::instrument(skip_all, fields(otel.name = "graphQL", otel.kind = ?SpanKind::Server))]
@@ -134,15 +127,10 @@ pub async fn graphql_request<T: DeserializeOwned + GraphQLRequestLike>(
     match graphql_request {
         Ok(request) => {
             let mut response = request.data(req_ctx.clone()).execute(&app_ctx.schema).await;
-            let cookie_headers = req_ctx.cookie_headers.clone();
+
             response = update_cache_control_header(response, app_ctx, req_ctx.clone());
             let mut resp = response.to_response()?;
-            update_response_headers(
-                &mut resp,
-                cookie_headers.map(|v| v.lock().unwrap().clone()),
-                app_ctx,
-            );
-            update_experimental_headers(&mut resp, app_ctx, req_ctx);
+            update_response_headers(&mut resp, &req_ctx, app_ctx);
             Ok(resp)
         }
         Err(err) => {
@@ -250,15 +238,9 @@ async fn handle_rest_apis(
                 .data(req_ctx.clone())
                 .execute(&app_ctx.schema)
                 .await;
-            let cookie_headers = req_ctx.cookie_headers.clone();
             response = update_cache_control_header(response, app_ctx.as_ref(), req_ctx.clone());
             let mut resp = response.to_rest_response()?;
-            update_response_headers(
-                &mut resp,
-                cookie_headers.map(|v| v.lock().unwrap().clone()),
-                app_ctx.as_ref(),
-            );
-            update_experimental_headers(&mut resp, app_ctx.as_ref(), req_ctx);
+            update_response_headers(&mut resp, &req_ctx, &app_ctx);
             Ok(resp)
         }
         .instrument(span)

src/lambda/evaluation_context.rs

@@ -78,7 +78,7 @@ impl<'a, Ctx: ResolverContextLike<'a>> EvaluationContext<'a, Ctx> {
     }
 
     pub fn headers(&self) -> &HeaderMap {
-        &self.req_ctx.req_headers
+        &self.req_ctx.request_headers
     }
 
     pub fn header(&self, key: &str) -> Option<&str> {

src/lambda/io.rs

@@ -157,6 +157,7 @@ fn set_headers<'ctx, Ctx: ResolverContextLike<'ctx>>(
 ) {
     set_cache_control(ctx, res);
     set_cookie_headers(ctx, res);
+    set_experimental_headers(ctx, res);
 }
 
 fn set_cache_control<'ctx, Ctx: ResolverContextLike<'ctx>>(
@@ -170,6 +171,13 @@ fn set_cache_control<'ctx, Ctx: ResolverContextLike<'ctx>>(
     }
 }
 
+fn set_experimental_headers<'ctx, Ctx: ResolverContextLike<'ctx>>(
+    ctx: &EvaluationContext<'ctx, Ctx>,
+    res: &Response<async_graphql::Value>,
+) {
+    ctx.req_ctx.add_x_headers(&res.headers);
+}
+
 fn set_cookie_headers<'ctx, Ctx: ResolverContextLike<'ctx>>(
     ctx: &EvaluationContext<'ctx, Ctx>,
     res: &Response<async_graphql::Value>,
@@ -183,13 +191,16 @@ async fn execute_raw_request<'ctx, Ctx: ResolverContextLike<'ctx>>(
     ctx: &EvaluationContext<'ctx, Ctx>,
     req: Request,
 ) -> Result<Response<async_graphql::Value>> {
-    ctx.req_ctx
+    let response = ctx
+        .req_ctx
         .runtime
         .http
         .execute(req)
         .await
         .map_err(|e| EvaluationError::IOException(e.to_string()))?
-        .to_json()
+        .to_json()?;
+
+    Ok(response)
 }
 
 async fn execute_raw_grpc_request<'ctx, Ctx: ResolverContextLike<'ctx>>(

src/path.rs

@@ -209,7 +209,7 @@ mod tests {
         }
 
         static REQ_CTX: Lazy<RequestContext> = Lazy::new(|| {
-            let mut req_ctx = RequestContext::default().req_headers(TEST_HEADERS.clone());
+            let mut req_ctx = RequestContext::default().request_headers(TEST_HEADERS.clone());
 
             req_ctx.server.vars = TEST_VARS.clone();
             req_ctx.runtime.env = Arc::new(Env::init(TEST_ENV_VARS.clone()));

tests/execution/experimental-headers.md

@@ -19,13 +19,13 @@ type User {
 - request:
     method: GET
     url: http://jsonplaceholder.typicode.com/users
+    body: null
+  response:
+    status: 200
     headers:
       X-tailcall: "tailcall-header"
       x-experimental: "experimental-header"
       x-not-allowed: "not-allowed-header"
-    body: null
-  response:
-    status: 200
     body:
       - id: 1
         name: Leanne Graham
@@ -34,9 +34,6 @@ type User {
 ```yml @assert
 - method: POST
   url: http://localhost:8080/graphql
-  headers:
-    X-tailcall: "tailcall-header"
-    x-experimental: "experimental-header"
   body:
     query: query { users { id name } }
 ```