synapse_proxy/transform/
mod.rs

1//! The transform pipeline: traits a request/response step implements, the
2//! mutable request/response views they operate on, and the registry of
3//! code-registered custom transforms.
4
5pub mod error_remap;
6pub mod inject;
7pub mod wrap;
8
9use std::collections::HashMap;
10use std::sync::Arc;
11
12use async_trait::async_trait;
13use axum::http::{HeaderMap, HeaderName, HeaderValue, Method, StatusCode};
14use serde_json::Value;
15
16use crate::context::ResolvedContext;
17
18/// A step's failure: `Reject` short-circuits the pipeline with a contract body.
19#[derive(Debug, Clone)]
20pub enum TransformError {
21    Reject {
22        status: StatusCode,
23        error: String,
24        detail: String,
25    },
26    Internal(String),
27}
28
29/// Mutable view of the request a `RequestTransform` may edit before forwarding.
30/// `body` is JSON when present; non-JSON bodies are exposed as raw bytes and are
31/// not editable as JSON.
32pub struct ProxyRequest {
33    pub method: Method,
34    pub path: String,
35    pub query: Option<String>,
36    pub headers: HeaderMap,
37    body: Body,
38}
39
40enum Body {
41    Bytes(Vec<u8>),
42    Json(Value),
43}
44
45impl ProxyRequest {
46    pub fn from_parts(
47        method: Method,
48        path: String,
49        query: Option<String>,
50        headers: HeaderMap,
51        bytes: Vec<u8>,
52    ) -> Self {
53        Self {
54            method,
55            path,
56            query,
57            headers,
58            body: Body::Bytes(bytes),
59        }
60    }
61
62    /// Overwrite a header (skips silently if name/value is invalid).
63    pub fn set_header(&mut self, name: &str, value: &str) {
64        if let (Ok(n), Ok(v)) = (HeaderName::try_from(name), HeaderValue::try_from(value)) {
65            self.headers.insert(n, v);
66        }
67    }
68
69    /// Remove a header entirely (used to strip a caller value when an injected
70    /// context key is absent — fail-safe identity handling).
71    pub fn remove_header(&mut self, name: &str) {
72        if let Ok(n) = HeaderName::try_from(name) {
73            self.headers.remove(n);
74        }
75    }
76
77    /// Mutable JSON body, parsing the raw bytes on first access. Errors if the
78    /// body is not valid JSON (a body transform on a non-JSON body is a reject).
79    pub fn body_json_mut(&mut self) -> Result<&mut Value, TransformError> {
80        if let Body::Bytes(b) = &self.body {
81            let parsed = if b.is_empty() {
82                Value::Object(Default::default())
83            } else {
84                serde_json::from_slice(b).map_err(|e| TransformError::Reject {
85                    status: StatusCode::BAD_REQUEST,
86                    error: "invalid_body".into(),
87                    detail: format!("expected JSON body: {e}"),
88                })?
89            };
90            self.body = Body::Json(parsed);
91        }
92        match &mut self.body {
93            Body::Json(v) => Ok(v),
94            Body::Bytes(_) => unreachable!(),
95        }
96    }
97
98    /// Serialize the (possibly transformed) body back to bytes for forwarding.
99    pub fn into_body_bytes(self) -> Vec<u8> {
100        match self.body {
101            Body::Bytes(b) => b,
102            Body::Json(v) => serde_json::to_vec(&v).unwrap_or_default(),
103        }
104    }
105
106    /// Headers + body ready for the upstream reqwest call.
107    ///
108    /// When JSON transforms (`inject`/`wrap` on body) re-serialize the payload,
109    /// drop stale `Content-Length` / `Transfer-Encoding` from the client so
110    /// reqwest emits a length matching the new body (callers like urllib set
111    /// Content-Length for the pre-transform size).
112    pub fn into_forward_parts(self) -> (HeaderMap, Vec<u8>) {
113        let reencoded = matches!(self.body, Body::Json(_));
114        let bytes = match self.body {
115            Body::Bytes(b) => b,
116            Body::Json(v) => serde_json::to_vec(&v).unwrap_or_default(),
117        };
118        let mut headers = self.headers;
119        if reencoded {
120            headers.remove("content-length");
121            headers.remove("transfer-encoding");
122        }
123        (headers, bytes)
124    }
125}
126
127/// Mutable view of the upstream response. v1 transforms touch status/headers and
128/// may replace the body with a small JSON value; if no replacement is set the
129/// handler streams the upstream body unchanged.
130pub struct ProxyResponse {
131    pub status: StatusCode,
132    pub headers: HeaderMap,
133    replacement: Option<Value>,
134}
135
136impl ProxyResponse {
137    pub fn new(status: StatusCode, headers: HeaderMap) -> Self {
138        Self {
139            status,
140            headers,
141            replacement: None,
142        }
143    }
144    /// Replace the response body with `value` (sent instead of streaming upstream).
145    pub fn replace_body(&mut self, value: Value) {
146        self.replacement = Some(value);
147    }
148    pub fn replacement(&self) -> Option<&Value> {
149        self.replacement.as_ref()
150    }
151}
152
153#[async_trait]
154pub trait RequestTransform: Send + Sync {
155    async fn apply(
156        &self,
157        ctx: &ResolvedContext,
158        req: &mut ProxyRequest,
159    ) -> Result<(), TransformError>;
160}
161
162#[async_trait]
163pub trait ResponseTransform: Send + Sync {
164    async fn apply(
165        &self,
166        ctx: &ResolvedContext,
167        resp: &mut ProxyResponse,
168    ) -> Result<(), TransformError>;
169}
170
171/// Code-registered custom transforms, looked up by `{ transform = "name" }`.
172#[derive(Default, Clone)]
173pub struct TransformRegistry {
174    pub(crate) request: HashMap<String, Arc<dyn RequestTransform>>,
175    pub(crate) response: HashMap<String, Arc<dyn ResponseTransform>>,
176}
177
178impl TransformRegistry {
179    pub fn register_request(&mut self, name: impl Into<String>, t: Arc<dyn RequestTransform>) {
180        self.request.insert(name.into(), t);
181    }
182    pub fn register_response(&mut self, name: impl Into<String>, t: Arc<dyn ResponseTransform>) {
183        self.response.insert(name.into(), t);
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    fn req() -> ProxyRequest {
192        ProxyRequest::from_parts(
193            Method::POST,
194            "/x".into(),
195            None,
196            HeaderMap::new(),
197            b"{\"a\":1}".to_vec(),
198        )
199    }
200
201    #[test]
202    fn set_header_overwrites() {
203        let mut r = req();
204        r.set_header("x-test", "v1");
205        r.set_header("x-test", "v2");
206        assert_eq!(r.headers.get("x-test").unwrap(), "v2");
207    }
208
209    #[test]
210    fn body_json_mut_parses_and_serializes() {
211        let mut r = req();
212        r.body_json_mut().unwrap()["b"] = serde_json::json!(2);
213        let bytes = r.into_body_bytes();
214        let v: Value = serde_json::from_slice(&bytes).unwrap();
215        assert_eq!(v["a"], 1);
216        assert_eq!(v["b"], 2);
217    }
218
219    #[test]
220    fn into_forward_parts_strips_stale_content_length_after_json_inject() {
221        let mut headers = HeaderMap::new();
222        headers.insert("content-length", "10".parse().unwrap());
223        headers.insert("transfer-encoding", "chunked".parse().unwrap());
224        let mut r = ProxyRequest::from_parts(
225            Method::POST,
226            "/x".into(),
227            None,
228            headers,
229            br#"{"a":1}"#.to_vec(),
230        );
231        r.body_json_mut().unwrap()["orgId"] = serde_json::json!("acme");
232        let (headers, bytes) = r.into_forward_parts();
233        assert!(headers.get("content-length").is_none());
234        assert!(headers.get("transfer-encoding").is_none());
235        let v: Value = serde_json::from_slice(&bytes).unwrap();
236        assert_eq!(v["orgId"], "acme");
237        assert!(bytes.len() > 10);
238    }
239
240    #[test]
241    fn into_forward_parts_keeps_content_length_when_body_unchanged() {
242        let mut headers = HeaderMap::new();
243        headers.insert("content-length", "7".parse().unwrap());
244        let r = ProxyRequest::from_parts(
245            Method::POST,
246            "/x".into(),
247            None,
248            headers,
249            br#"{"a":1}"#.to_vec(),
250        );
251        let (headers, bytes) = r.into_forward_parts();
252        assert_eq!(headers.get("content-length").unwrap(), "7");
253        assert_eq!(bytes, br#"{"a":1}"#);
254    }
255}
synapse_proxy/transform/mod.rs

synapse_proxy/transform/
mod.rs
