Apollo Router forward subgraph headers to clients

Help
,
1

I have a node gateway right now and I was looking to move to Apollo Router. The only functionality that doesn’t come out of the box for Router that I need to port over from my gateway is to return headers to the client that were set by a subgraph. I took a stab at a custom plugin for this but the headers set by my subgraphs are still not being passed through to the client. Here’s what my plugin looks like:

use apollo_router_core::{
    register_plugin, Plugin, RouterRequest, RouterResponse, SubgraphResponse, SubgraphRequest,
};
use http::HeaderValue;
use tower::{util::BoxService, BoxError, ServiceBuilder, ServiceExt};
use cookie::{Cookie, SameSite};
use schemars::JsonSchema;
use serde::Deserialize;
use once_cell::sync::OnceCell;

static SET_COOKIE_HEADER: &str = "set-cookie";
static DOMAIN: OnceCell<String> = OnceCell::new();

#[derive(Debug, Default, Deserialize, JsonSchema)]
struct Conf {
    allow_local_host: bool,
    cookie_domain: String,
    cookie_same_site: String,
    use_secure_cookie: bool,
}

// We are storing the configuration, but not using it. Hence the allow dead code.
#[derive(Debug, Default)]
struct CookieSetter {
    allow_local_host: bool,
    cookie_same_site: Option<SameSite>,
    use_secure_cookie: bool,
}


#[async_trait::async_trait]
impl Plugin for CookieSetter {
    type Config = Conf;

    async fn new(configuration: Self::Config) -> Result<Self, BoxError> {
        let same_site = configuration.cookie_same_site.clone().to_lowercase();
        let cookie_same_site = match same_site.as_str() {
            "strict" => SameSite::Strict,
            "lax" => SameSite::Lax,
            _ => SameSite::None,
        };

        let domain = configuration.cookie_domain.clone();
        DOMAIN.set(domain).unwrap();
        return Ok(Self {
            allow_local_host: configuration.allow_local_host,
            cookie_same_site: Some(cookie_same_site),
            use_secure_cookie: configuration.use_secure_cookie,
        })
    }

    fn router_service(
        &mut self,
        service: BoxService<RouterRequest, RouterResponse, BoxError>,
    ) -> BoxService<RouterRequest, RouterResponse, BoxError> {
        // `ServiceBuilder` provides us with `map_request` and `map_response` methods.
        //
        // These allow basic interception and transformation of request and response messages.
        let same_site = self.cookie_same_site.clone();
        let allow_local_host = self.allow_local_host.clone();
        let use_secure_cookie = self.use_secure_cookie.clone();

        ServiceBuilder::new()
            .service(service)
            .map_response(move |mut resp: RouterResponse| {
                let headers = resp.response.headers_mut();
                for (key, val) in headers.clone() {
                    match key {
                        Some(name) => {
                            let name_lower = name.as_str().to_lowercase();
                            if name_lower == SET_COOKIE_HEADER {
                                let cookie_string = val.to_str().unwrap().to_owned();
                                let result = Cookie::parse(cookie_string);
                                match result {
                                    Ok(mut cookie) => {
                                        cookie.set_same_site(same_site);
                                        cookie.set_secure(use_secure_cookie);
                                        if allow_local_host {
                                            cookie.unset_domain();
                                        } else {
                                            cookie.set_domain(DOMAIN.get().unwrap());
                                        }
                                        let new_header = cookie.to_string();
                                        let header_value = HeaderValue::from_str(new_header.as_str()).unwrap();
                                        headers.insert(name, header_value);
                                    },
                                    Err(_) => (),
                                }
                            }
                        },
                        None => (),
                    }
                }
                resp
            })
            .boxed()
    }

    fn subgraph_service(
        &mut self,
        _nam: &str,
        service: BoxService<SubgraphRequest, SubgraphResponse, BoxError>,
    ) -> BoxService<SubgraphRequest, SubgraphResponse, BoxError> {
        // `ServiceBuilder` provides us with `map_request` and `map_response` methods.
        //
        // These allow basic interception and transformation of request and response messages.
        let same_site = self.cookie_same_site.clone();
        let allow_local_host = self.allow_local_host.clone();
        let use_secure_cookie = self.use_secure_cookie.clone();

        ServiceBuilder::new()
            .service(service)
            .map_response(move |mut resp: SubgraphResponse| {
                let headers = resp.response.headers_mut();
                for (key, val) in headers.clone() {
                    match key {
                        Some(name) => {
                            let name_lower = name.as_str().to_lowercase();
                            if name_lower == SET_COOKIE_HEADER {
                                let cookie_string = val.to_str().unwrap().to_owned();
                                let result = Cookie::parse(cookie_string);
                                match result {
                                    Ok(mut cookie) => {
                                        cookie.set_same_site(same_site);
                                        cookie.set_secure(use_secure_cookie);
                                        if allow_local_host {
                                            cookie.unset_domain();
                                        } else {
                                            cookie.set_domain(DOMAIN.get().unwrap());
                                        }
                                        let new_header = cookie.to_string();
                                        let header_value = HeaderValue::from_str(new_header.as_str()).unwrap();
                                        headers.insert(name, header_value);
                                    },
                                    Err(_) => (),
                                }
                            }
                        },
                        None => (),
                    }
                }
                resp
            })
            .boxed()
    }
}

register_plugin!("plugins", "cookie", CookieSetter);
2

Your plugin looks great! Thanks for sharing it with us.

The reason it isn’t giving you the results you want is that you are manipulating the response at the Subgraph and then assuming, reasonably, that the response you get in the Router is the same response you previously manipulated, which it isn’t.

Explanation: At the subgraph level, you will be getting responses from at least one subgraph, often many subgraphs, and so processing returns from them is complex in terms of both sequencing and differentiating responses from different subgraphs. The router pipeline reports various response details for stage appropriate processing, but the actual subgraph response is not propagated back through the whole pipeline. This is why the manipulation of headers that you are doing in your subgraph response processing is “lost”. The pipeline discards the response at the Execution stage and builds a new response which is ultimately returned to the client.

(Note: We are aware that this is problematic for users and there are tickets filed to streamline and improve the pipeline process which will make this kind of task simpler.)

What you can do right now is use the “Context” which we provide for storing arbitrary data. There is an example of how to use context from a plugin.

In fact, it may be possible (and easier) to do this with a Rhai script. We have another example which illustrates how to build a surrogate cache key which accesses header values returned by subgraphs. In the router response, we use the context-stored map of responses to generate a derived surrogate key. You could do something similar to generate either multiple set-cookie headers or merge them or …

I hope this is helpful. If not, feel free to follow up by logging issues against the router.

3

Thanks for the recommendations. I’ll update my plugin to use context to store the data I need during the subgraph map_response and set the headers from context during the router map_response. It would be really helpful for the subgraph map_response to also have access to the router response similar to express so that we can directly affect changes like this.