use std::sync::Arc;
use std::sync::Mutex;
use connectrpc::client::{BoxFuture, ClientBody, ClientTransport, HttpClient};
use connectrpc::compression::CompressionRegistry;
use connectrpc::ConnectError;
use http::header::{ACCEPT_ENCODING, COOKIE, SET_COOKIE};
use http::{Request, Response};
#[must_use]
pub fn connect_compression_registry() -> CompressionRegistry {
CompressionRegistry::default()
}
pub const EXOWARE_AFFINITY_COOKIE: &str = "exoware_affinity_cookie";
#[derive(Clone, Debug)]
pub struct PreferZstdHttpClient {
inner: HttpClient,
sticky_cookie: Arc<Mutex<Option<String>>>,
}
impl PreferZstdHttpClient {
pub fn plaintext() -> Self {
Self {
inner: HttpClient::plaintext(),
sticky_cookie: Arc::new(Mutex::new(None)),
}
}
}
impl ClientTransport for PreferZstdHttpClient {
type ResponseBody = hyper::body::Incoming;
type Error = ConnectError;
fn send(
&self,
mut request: Request<ClientBody>,
) -> BoxFuture<'static, Result<Response<Self::ResponseBody>, Self::Error>> {
if let Ok(guard) = self.sticky_cookie.lock() {
if let Some(ref pair) = *guard {
if let Ok(hv) = http::HeaderValue::from_str(pair) {
request.headers_mut().insert(COOKIE, hv);
}
}
}
request.headers_mut().insert(
ACCEPT_ENCODING,
http::HeaderValue::from_static("zstd, gzip"),
);
let inner = self.inner.clone();
let sticky_cookie = Arc::clone(&self.sticky_cookie);
Box::pin(async move {
let response = inner.send(request).await?;
let (parts, body) = response.into_parts();
if let Ok(mut g) = sticky_cookie.lock() {
for val in parts.headers.get_all(SET_COOKIE) {
if let Ok(s) = val.to_str() {
if let Some(pair) = parse_sticky_cookie_pair(s, EXOWARE_AFFINITY_COOKIE) {
*g = Some(pair);
break;
}
}
}
}
Ok(Response::from_parts(parts, body))
})
}
}
fn parse_sticky_cookie_pair(set_cookie: &str, name: &str) -> Option<String> {
let first = set_cookie.split(';').next()?.trim();
let rest = first.strip_prefix(name)?.strip_prefix('=')?;
let val = rest.trim().trim_matches('"');
if val.is_empty() {
return None;
}
Some(format!("{name}={val}"))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_sticky_cookie_pair_handles_quoted_value() {
let s = r#"exoware_affinity_cookie="Cg4xMjcuMC4wLjE6ODA4MQ=="; Path=/; HttpOnly"#;
assert_eq!(
parse_sticky_cookie_pair(s, EXOWARE_AFFINITY_COOKIE),
Some("exoware_affinity_cookie=Cg4xMjcuMC4wLjE6ODA4MQ==".to_string())
);
}
}