Skip to main content

exoware_sdk/proto/
compression.rs

1//! Wire compression for the store API.
2//!
3//! ## Registry
4//!
5//! Servers register **gzip** and **zstd** via [`connect_compression_registry`] (same as
6//! [`connectrpc::compression::CompressionRegistry::default`]) so callers without zstd
7//! (including typical browsers) can still negotiate gzip.
8//!
9//! ## Rust client transport
10//!
11//! HTTP transport that sets `Accept-Encoding: zstd, gzip` on every outbound request.
12//!
13//! [`connectrpc::compression::CompressionRegistry::default`] builds the header value in sorted
14//! order (`gzip, zstd`), so servers negotiate **gzip** first. Replacing the header after
15//! connectrpc builds the request lets clients **prefer zstd** while still advertising gzip.
16//!
17//! **Request bodies** (client -> server) use a single codec from connectrpc `compress_requests`.
18//!
19//! ## Edge upstream affinity (cookie)
20//!
21//! Deployments behind a load balancer or proxy may use HTTP sticky sessions: the edge sets
22//! `Set-Cookie` for [`EXOWARE_AFFINITY_COOKIE`] so each client session sticks to one backend
23//! (cache locality). This repo's Docker/Envoy example uses the stateful session filter for that.
24//! [`PreferZstdHttpClient`] stores `Set-Cookie` from responses and sends `Cookie` on subsequent RPCs.
25
26use std::sync::Arc;
27use std::sync::Mutex;
28
29use connectrpc::client::{BoxFuture, ClientBody, ClientTransport, HttpClient};
30use connectrpc::compression::CompressionRegistry;
31use connectrpc::ConnectError;
32use http::header::{ACCEPT_ENCODING, COOKIE, SET_COOKIE};
33use http::{Request, Response};
34
35/// gzip + zstd - used for [`connectrpc::ConnectRpcService::with_compression`] and
36/// [`connectrpc::client::ClientConfig::compression`].
37#[must_use]
38pub fn connect_compression_registry() -> CompressionRegistry {
39    CompressionRegistry::default()
40}
41
42/// Sticky-session cookie name; must match whatever the edge emits in `Set-Cookie`.
43pub const EXOWARE_AFFINITY_COOKIE: &str = "exoware_affinity_cookie";
44
45/// Wraps [`HttpClient`] so every RPC sends `Accept-Encoding: zstd, gzip` (see module docs).
46///
47/// Also persists **HTTP sticky-session** behavior for [`EXOWARE_AFFINITY_COOKIE`]: when responses
48/// include `Set-Cookie: exoware_affinity_cookie=...`, the value is stored and sent on later requests as
49/// `Cookie: exoware_affinity_cookie=...` so the same client handle stays pinned to one upstream.
50#[derive(Clone, Debug)]
51pub struct PreferZstdHttpClient {
52    inner: HttpClient,
53    /// `Cookie` header line body (`name=value`) for [`EXOWARE_AFFINITY_COOKIE`], no `Cookie:` prefix.
54    sticky_cookie: Arc<Mutex<Option<String>>>,
55}
56
57impl PreferZstdHttpClient {
58    pub fn plaintext() -> Self {
59        Self {
60            inner: HttpClient::plaintext(),
61            sticky_cookie: Arc::new(Mutex::new(None)),
62        }
63    }
64}
65
66impl ClientTransport for PreferZstdHttpClient {
67    type ResponseBody = hyper::body::Incoming;
68    type Error = ConnectError;
69
70    fn send(
71        &self,
72        mut request: Request<ClientBody>,
73    ) -> BoxFuture<'static, Result<Response<Self::ResponseBody>, Self::Error>> {
74        if let Ok(guard) = self.sticky_cookie.lock() {
75            if let Some(ref pair) = *guard {
76                if let Ok(hv) = http::HeaderValue::from_str(pair) {
77                    request.headers_mut().insert(COOKIE, hv);
78                }
79            }
80        }
81        request.headers_mut().insert(
82            ACCEPT_ENCODING,
83            http::HeaderValue::from_static("zstd, gzip"),
84        );
85        let inner = self.inner.clone();
86        let sticky_cookie = Arc::clone(&self.sticky_cookie);
87        Box::pin(async move {
88            let response = inner.send(request).await?;
89            let (parts, body) = response.into_parts();
90            if let Ok(mut g) = sticky_cookie.lock() {
91                for val in parts.headers.get_all(SET_COOKIE) {
92                    if let Ok(s) = val.to_str() {
93                        if let Some(pair) = parse_sticky_cookie_pair(s, EXOWARE_AFFINITY_COOKIE) {
94                            *g = Some(pair);
95                            break;
96                        }
97                    }
98                }
99            }
100            Ok(Response::from_parts(parts, body))
101        })
102    }
103}
104
105/// From one `Set-Cookie` header value, extract `name=value` for the affinity cookie.
106fn parse_sticky_cookie_pair(set_cookie: &str, name: &str) -> Option<String> {
107    let first = set_cookie.split(';').next()?.trim();
108    let rest = first.strip_prefix(name)?.strip_prefix('=')?;
109    let val = rest.trim().trim_matches('"');
110    if val.is_empty() {
111        return None;
112    }
113    Some(format!("{name}={val}"))
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119
120    #[test]
121    fn parse_sticky_cookie_pair_handles_quoted_value() {
122        let s = r#"exoware_affinity_cookie="Cg4xMjcuMC4wLjE6ODA4MQ=="; Path=/; HttpOnly"#;
123        assert_eq!(
124            parse_sticky_cookie_pair(s, EXOWARE_AFFINITY_COOKIE),
125            Some("exoware_affinity_cookie=Cg4xMjcuMC4wLjE6ODA4MQ==".to_string())
126        );
127    }
128}