1use async_trait::async_trait;
20use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
21use bytes::Bytes;
22use regex::Regex;
23use std::collections::HashMap;
24
25use crate::context::Context;
26use crate::plugins::util::content_codec::{self, ContentEncoding};
27use crate::plugins::{Plugin, PluginOutput, PluginResult};
28use crate::vars;
29
30pub struct ResponseRewritePlugin {
37 status_code: Option<u16>,
38 body: Option<Bytes>,
40 filters: Vec<BodyFilter>,
41 add_headers: Vec<(String, String)>,
42 set_headers: Vec<(String, String)>,
43 remove_headers: Vec<String>,
44 vars: Option<vars::Expr>,
45}
46
47struct BodyFilter {
49 regex: Regex,
50 replace: String,
51 global: bool,
52}
53
54const BODY_DERIVED_HEADERS: [&str; 4] = [
59 "content-length",
60 "content-encoding",
61 "last-modified",
62 "etag",
63];
64
65fn clear_body_derived_headers(ctx: &mut Context) {
66 for name in BODY_DERIVED_HEADERS {
67 crate::plugins::util::headers::remove_ci(&mut ctx.response.headers, name);
68 }
69}
70
71fn header_value_to_string(v: &serde_json::Value) -> Option<String> {
74 match v {
75 serde_json::Value::String(s) => Some(s.clone()),
76 serde_json::Value::Number(n) => Some(n.to_string()),
77 _ => None,
78 }
79}
80
81fn parse_add_entry(entry: &str) -> Result<(String, String), String> {
86 let (name, value) = entry
87 .split_once(':')
88 .ok_or_else(|| format!("headers.add entry '{}' must be 'Name: value'", entry))?;
89 let name = name.trim();
90 let value = value.trim();
91 if name.is_empty() || name.contains(char::is_whitespace) {
92 return Err(format!(
93 "headers.add entry '{}' has an invalid header name",
94 entry
95 ));
96 }
97 if value.is_empty() || value.contains(':') {
98 return Err(format!(
99 "headers.add entry '{}' has an invalid header value (must be non-empty, no ':')",
100 entry
101 ));
102 }
103 Ok((name.to_lowercase(), value.to_string()))
104}
105
106#[allow(clippy::type_complexity)]
110fn parse_headers(
111 raw: &serde_json::Value,
112) -> Result<(Vec<(String, String)>, Vec<(String, String)>, Vec<String>), String> {
113 let obj = raw
114 .as_object()
115 .ok_or("headers must be an object".to_string())?;
116
117 let is_structured = obj.get("add").is_some_and(|v| v.is_array())
118 || obj.get("set").is_some_and(|v| v.is_object())
119 || obj.get("remove").is_some_and(|v| v.is_array());
120
121 let mut add = Vec::new();
122 let mut set = Vec::new();
123 let mut remove = Vec::new();
124
125 if is_structured {
126 if let Some(entries) = obj.get("add") {
127 for entry in entries.as_array().ok_or("headers.add must be an array")? {
128 let s = entry
129 .as_str()
130 .ok_or("headers.add entries must be strings ('Name: value')")?;
131 add.push(parse_add_entry(s)?);
132 }
133 }
134 if let Some(map) = obj.get("set") {
135 for (name, value) in map.as_object().ok_or("headers.set must be a map")? {
136 let v = header_value_to_string(value)
137 .ok_or_else(|| format!("headers.set['{}'] must be a string or number", name))?;
138 set.push((name.to_lowercase(), v));
139 }
140 }
141 if let Some(names) = obj.get("remove") {
142 for name in names.as_array().ok_or("headers.remove must be an array")? {
143 let s = name
144 .as_str()
145 .ok_or("headers.remove entries must be strings")?;
146 remove.push(s.to_lowercase());
147 }
148 }
149 } else {
150 for (name, value) in obj {
152 let v = header_value_to_string(value)
153 .ok_or_else(|| format!("headers['{}'] must be a string or number", name))?;
154 set.push((name.to_lowercase(), v));
155 }
156 }
157
158 Ok((add, set, remove))
159}
160
161fn parse_filter(v: &serde_json::Value) -> Result<BodyFilter, String> {
163 let obj = v
164 .as_object()
165 .ok_or("filters entries must be objects".to_string())?;
166
167 let pattern = obj
168 .get("regex")
169 .and_then(|v| v.as_str())
170 .filter(|s| !s.is_empty())
171 .ok_or("filters entries require a non-empty 'regex' string")?;
172
173 let replace = obj
174 .get("replace")
175 .and_then(|v| v.as_str())
176 .ok_or("filters entries require a 'replace' string")?
177 .to_string();
178
179 let global = match obj.get("scope").and_then(|v| v.as_str()) {
180 None | Some("once") => false,
181 Some("global") => true,
182 Some(other) => {
183 return Err(format!(
184 "filters scope must be 'once' or 'global', got '{}'",
185 other
186 ))
187 }
188 };
189
190 let options = obj.get("options").and_then(|v| v.as_str()).unwrap_or("");
191 let case_insensitive = match options {
192 "" => false,
193 "i" => true,
194 other => {
195 return Err(format!(
196 "filters options only supports 'i' (case-insensitive), got '{}'",
197 other
198 ))
199 }
200 };
201
202 let full_pattern = if case_insensitive {
203 format!("(?i){}", pattern)
204 } else {
205 pattern.to_string()
206 };
207 let regex = Regex::new(&full_pattern)
208 .map_err(|e| format!("filters regex \"{}\" validation failed: {}", pattern, e))?;
209
210 Ok(BodyFilter {
211 regex,
212 replace,
213 global,
214 })
215}
216
217impl ResponseRewritePlugin {
218 pub fn from_config(config: &HashMap<String, serde_json::Value>) -> Result<Self, String> {
264 for wrong in ["add_headers", "set_headers", "remove_headers"] {
269 if config.contains_key(wrong) {
270 return Err(format!(
271 "response-rewrite has no '{wrong}' key (that is proxy-rewrite's schema); \
272 use headers with add/set/remove, e.g. \
273 headers: {{ set: {{ x-foo: bar }}, remove: [x-powered-by] }}"
274 ));
275 }
276 }
277
278 let status_code = match config.get("status_code") {
279 None => None,
280 Some(v) => {
281 let code = v
282 .as_u64()
283 .ok_or("status_code must be an integer".to_string())?;
284 if !(200..=598).contains(&code) {
285 return Err(format!("status_code must be within 200-598, got {}", code));
286 }
287 Some(code as u16)
288 }
289 };
290
291 let body_base64 = config
292 .get("body_base64")
293 .and_then(|v| v.as_bool())
294 .unwrap_or(false);
295
296 let body = match config.get("body") {
297 None => {
298 if body_base64 {
299 return Err("body_base64 requires 'body' to be set".to_string());
300 }
301 None
302 }
303 Some(v) => {
304 let s = v.as_str().ok_or("body must be a string".to_string())?;
305 if body_base64 {
306 if s.is_empty() {
307 return Err("invalid base64 content".to_string());
308 }
309 let decoded = BASE64
310 .decode(s.trim())
311 .map_err(|_| "invalid base64 content".to_string())?;
312 Some(Bytes::from(decoded))
313 } else {
314 Some(Bytes::from(s.to_string()))
315 }
316 }
317 };
318
319 let filters = match config.get("filters") {
320 None => Vec::new(),
321 Some(v) => {
322 let arr = v.as_array().ok_or("filters must be an array".to_string())?;
323 if arr.is_empty() {
324 return Err("filters must contain at least one entry".to_string());
325 }
326 arr.iter()
327 .map(parse_filter)
328 .collect::<Result<Vec<_>, _>>()?
329 }
330 };
331
332 if body.is_some() && !filters.is_empty() {
333 return Err("'body' and 'filters' are mutually exclusive".to_string());
334 }
335
336 let (add_headers, set_headers, remove_headers) = match config.get("headers") {
337 None => (Vec::new(), Vec::new(), Vec::new()),
338 Some(raw) => parse_headers(raw)?,
339 };
340
341 let vars = match config.get("vars") {
342 None => None,
343 Some(v) => Some(
344 vars::Expr::parse(v)
345 .map_err(|e| format!("failed to validate the 'vars' expression: {}", e))?,
346 ),
347 };
348
349 Ok(Self {
350 status_code,
351 body,
352 filters,
353 add_headers,
354 set_headers,
355 remove_headers,
356 vars,
357 })
358 }
359
360 fn apply_filters(&self, ctx: &mut Context) {
366 let encoding_header = ctx
367 .response
368 .headers
369 .get("content-encoding")
370 .and_then(|v| v.first())
371 .cloned()
372 .unwrap_or_default();
373
374 let decoded = match ContentEncoding::parse(&encoding_header) {
375 Err(e) => {
376 tracing::warn!(
377 "response-rewrite: filters skipped due to unsupported \
378 compression encoding: {}",
379 e
380 );
381 return;
382 }
383 Ok(None) => ctx.response.body.clone(),
384 Ok(Some(encoding)) => match content_codec::decode(&encoding, &ctx.response.body) {
385 Ok(decoded) => decoded,
386 Err(e) => {
387 tracing::warn!("response-rewrite: filters skipped: {}", e);
388 return;
389 }
390 },
391 };
392
393 let mut text = match std::str::from_utf8(&decoded) {
394 Ok(s) => s.to_string(),
395 Err(_) => {
396 tracing::warn!("response-rewrite: filters skipped: body is not valid UTF-8");
397 return;
398 }
399 };
400
401 for filter in &self.filters {
402 text = if filter.global {
403 filter.regex.replace_all(&text, filter.replace.as_str())
404 } else {
405 filter.regex.replace(&text, filter.replace.as_str())
406 }
407 .into_owned();
408 }
409
410 ctx.response.body = Bytes::from(text);
413 clear_body_derived_headers(ctx);
414 }
415}
416
417#[async_trait]
418impl Plugin for ResponseRewritePlugin {
419 fn plugin_type(&self) -> &str {
420 "response-rewrite"
421 }
422
423 async fn execute(
424 &self,
425 mut ctx: Context,
426 _named_inputs: &HashMap<String, serde_json::Value>,
427 ) -> PluginResult {
428 if let Some(expr) = &self.vars {
430 if !expr.eval(&ctx) {
431 return Ok(PluginOutput {
432 context: ctx,
433 named_outputs: HashMap::new(),
434 });
435 }
436 }
437
438 if let Some(code) = self.status_code {
439 ctx.response.status_code = code;
440 }
441
442 if let Some(body) = &self.body {
443 ctx.response.body = body.clone();
444 clear_body_derived_headers(&mut ctx);
445 } else if !self.filters.is_empty() {
446 self.apply_filters(&mut ctx);
447 }
448
449 for (name, value) in &self.add_headers {
452 let value = vars::interpolate(&ctx, value);
453 ctx.response
454 .headers
455 .entry(name.clone())
456 .or_default()
457 .push(value);
458 }
459 for (name, value) in &self.set_headers {
460 let value = vars::interpolate(&ctx, value);
461 ctx.response.headers.insert(name.clone(), vec![value]);
462 }
463 for name in &self.remove_headers {
464 crate::plugins::util::headers::remove_ci(&mut ctx.response.headers, name);
465 }
466
467 Ok(PluginOutput {
468 context: ctx,
469 named_outputs: HashMap::new(),
470 })
471 }
472}
473
474#[cfg(test)]
475mod tests {
476 use super::*;
477 use crate::context::{GatewayRequest, GatewayResponse, Protocol};
478
479 fn test_context(status: u16, body: &[u8]) -> Context {
480 let mut response_headers = HashMap::new();
481 response_headers.insert("content-length".to_string(), vec![body.len().to_string()]);
482 Context {
483 request: GatewayRequest {
484 method: "GET".to_string(),
485 path: "/test".to_string(),
486 host: "localhost".to_string(),
487 scheme: "http".to_string(),
488 headers: HashMap::new(),
489 query_params: HashMap::new(),
490 body: Bytes::new(),
491 remote_addr: "127.0.0.1:12345".to_string(),
492 protocol: Protocol::Http1,
493 },
494 response: GatewayResponse {
495 status_code: status,
496 headers: response_headers,
497 body: Bytes::copy_from_slice(body),
498 },
499 message: HashMap::new(),
500 errors: Vec::new(),
501 }
502 }
503
504 fn plugin(config: serde_json::Value) -> ResponseRewritePlugin {
505 let map: HashMap<String, serde_json::Value> =
506 serde_json::from_value(config).expect("test config must be an object");
507 ResponseRewritePlugin::from_config(&map).expect("config should be valid")
508 }
509
510 #[tokio::test]
511 async fn test_response_rewrite_status_and_body() {
512 let p = plugin(serde_json::json!({
513 "status_code": 404,
514 "body": "not found\n"
515 }));
516 let ctx = test_context(200, b"original");
517 let out = p.execute(ctx, &HashMap::new()).await.unwrap();
518 assert_eq!(out.context.response.status_code, 404);
519 assert_eq!(out.context.response.body.as_ref(), b"not found\n");
520 assert!(!out.context.response.headers.contains_key("content-length"));
522 }
523
524 #[tokio::test]
525 async fn test_response_rewrite_body_base64() {
526 let p = plugin(serde_json::json!({
527 "body": "aGVsbG8gd29ybGQ=",
528 "body_base64": true
529 }));
530 let ctx = test_context(200, b"x");
531 let out = p.execute(ctx, &HashMap::new()).await.unwrap();
532 assert_eq!(out.context.response.body.as_ref(), b"hello world");
533 }
534
535 #[test]
536 fn test_response_rewrite_invalid_base64_rejected() {
537 let mut config = HashMap::new();
538 config.insert("body".to_string(), serde_json::json!("not!!valid@@base64"));
539 config.insert("body_base64".to_string(), serde_json::json!(true));
540 assert!(ResponseRewritePlugin::from_config(&config).is_err());
541
542 let mut config = HashMap::new();
544 config.insert("body_base64".to_string(), serde_json::json!(true));
545 assert!(ResponseRewritePlugin::from_config(&config).is_err());
546 }
547
548 #[test]
549 fn test_response_rewrite_config_validation() {
550 let mut config = HashMap::new();
552 config.insert("status_code".to_string(), serde_json::json!(199));
553 assert!(ResponseRewritePlugin::from_config(&config).is_err());
554 let mut config = HashMap::new();
555 config.insert("status_code".to_string(), serde_json::json!(599));
556 assert!(ResponseRewritePlugin::from_config(&config).is_err());
557
558 let mut config = HashMap::new();
560 config.insert("body".to_string(), serde_json::json!("x"));
561 config.insert(
562 "filters".to_string(),
563 serde_json::json!([{ "regex": "a", "replace": "b" }]),
564 );
565 assert!(ResponseRewritePlugin::from_config(&config).is_err());
566
567 let mut config = HashMap::new();
569 config.insert(
570 "filters".to_string(),
571 serde_json::json!([{ "regex": "(", "replace": "" }]),
572 );
573 assert!(ResponseRewritePlugin::from_config(&config).is_err());
574
575 let mut config = HashMap::new();
577 config.insert(
578 "filters".to_string(),
579 serde_json::json!([{ "regex": "a", "replace": "b", "options": "jo" }]),
580 );
581 assert!(ResponseRewritePlugin::from_config(&config).is_err());
582
583 let mut config = HashMap::new();
585 config.insert(
586 "headers".to_string(),
587 serde_json::json!({ "add": ["x-key: a:b"] }),
588 );
589 assert!(ResponseRewritePlugin::from_config(&config).is_err());
590 }
591
592 #[test]
595 fn test_proxy_rewrite_header_keys_rejected_with_guidance() {
596 for wrong in ["add_headers", "set_headers", "remove_headers"] {
597 let mut config = HashMap::new();
598 config.insert(wrong.to_string(), serde_json::json!({ "x-foo": "bar" }));
599 let err = ResponseRewritePlugin::from_config(&config)
600 .err()
601 .unwrap_or_else(|| panic!("{wrong} should be rejected"));
602 assert!(err.contains(wrong), "error should name the bad key: {err}");
603 assert!(
604 err.contains("headers"),
605 "error should point to the right shape: {err}"
606 );
607 }
608 let mut ok = HashMap::new();
610 ok.insert(
611 "headers".to_string(),
612 serde_json::json!({ "set": { "x-foo": "bar" } }),
613 );
614 assert!(ResponseRewritePlugin::from_config(&ok).is_ok());
615 }
616
617 #[tokio::test]
618 async fn test_response_rewrite_headers_add_set_remove() {
619 let p = plugin(serde_json::json!({
620 "headers": {
621 "add": ["X-Trace: abc"],
622 "set": { "X-Server": "featherbit", "x-version": 3 },
623 "remove": ["X-Powered-By"]
624 }
625 }));
626 let mut ctx = test_context(200, b"body");
627 ctx.response
628 .headers
629 .insert("x-powered-by".to_string(), vec!["php".to_string()]);
630 ctx.response
631 .headers
632 .insert("x-trace".to_string(), vec!["existing".to_string()]);
633 ctx.response
634 .headers
635 .insert("x-server".to_string(), vec!["nginx".to_string()]);
636
637 let out = p.execute(ctx, &HashMap::new()).await.unwrap();
638 let headers = &out.context.response.headers;
639 assert_eq!(
641 headers.get("x-trace"),
642 Some(&vec!["existing".to_string(), "abc".to_string()])
643 );
644 assert_eq!(
646 headers.get("x-server"),
647 Some(&vec!["featherbit".to_string()])
648 );
649 assert_eq!(headers.get("x-version"), Some(&vec!["3".to_string()]));
650 assert!(!headers.contains_key("x-powered-by"));
652 assert!(headers.contains_key("content-length"));
654 }
655
656 #[tokio::test]
657 async fn test_response_rewrite_deprecated_flat_headers_map() {
658 let p = plugin(serde_json::json!({
659 "headers": { "X-Flat": "yes" }
660 }));
661 let ctx = test_context(200, b"body");
662 let out = p.execute(ctx, &HashMap::new()).await.unwrap();
663 assert_eq!(
664 out.context.response.headers.get("x-flat"),
665 Some(&vec!["yes".to_string()])
666 );
667 }
668
669 #[tokio::test]
670 async fn test_response_rewrite_header_var_interpolation() {
671 let p = plugin(serde_json::json!({
672 "headers": {
673 "set": { "x-origin": "$remote_addr", "x-status": "$status" }
674 }
675 }));
676 let ctx = test_context(201, b"body");
677 let out = p.execute(ctx, &HashMap::new()).await.unwrap();
678 assert_eq!(
679 out.context.response.headers.get("x-origin"),
680 Some(&vec!["127.0.0.1".to_string()])
681 );
682 assert_eq!(
683 out.context.response.headers.get("x-status"),
684 Some(&vec!["201".to_string()])
685 );
686 }
687
688 #[tokio::test]
689 async fn test_response_rewrite_filters_once_and_global() {
690 let p = plugin(serde_json::json!({
691 "filters": [{ "regex": "foo", "replace": "bar" }]
692 }));
693 let ctx = test_context(200, b"foo foo foo");
694 let out = p.execute(ctx, &HashMap::new()).await.unwrap();
695 assert_eq!(out.context.response.body.as_ref(), b"bar foo foo");
696 assert!(!out.context.response.headers.contains_key("content-length"));
697
698 let p = plugin(serde_json::json!({
699 "filters": [{ "regex": "FOO", "replace": "bar", "scope": "global", "options": "i" }]
700 }));
701 let ctx = test_context(200, b"foo Foo fOO");
702 let out = p.execute(ctx, &HashMap::new()).await.unwrap();
703 assert_eq!(out.context.response.body.as_ref(), b"bar bar bar");
704 }
705
706 #[tokio::test]
707 async fn test_response_rewrite_filters_decode_gzip_body() {
708 let plain = Bytes::from_static(b"hello encoded world");
709 let compressed = content_codec::encode(&ContentEncoding::Gzip, &plain, 6).unwrap();
710
711 let p = plugin(serde_json::json!({
712 "filters": [{ "regex": "encoded", "replace": "decoded" }]
713 }));
714 let mut ctx = test_context(200, &compressed);
715 ctx.response
716 .headers
717 .insert("content-encoding".to_string(), vec!["gzip".to_string()]);
718 ctx.response
719 .headers
720 .insert("etag".to_string(), vec!["\"abc\"".to_string()]);
721
722 let out = p.execute(ctx, &HashMap::new()).await.unwrap();
723 assert_eq!(out.context.response.body.as_ref(), b"hello decoded world");
725 let headers = &out.context.response.headers;
726 assert!(!headers.contains_key("content-encoding"));
727 assert!(!headers.contains_key("content-length"));
728 assert!(!headers.contains_key("etag"));
729 }
730
731 #[tokio::test]
732 async fn test_response_rewrite_filters_skip_unsupported_encoding() {
733 let p = plugin(serde_json::json!({
734 "filters": [{ "regex": "foo", "replace": "bar" }]
735 }));
736 let mut ctx = test_context(200, b"foo body");
737 ctx.response
738 .headers
739 .insert("content-encoding".to_string(), vec!["zstd".to_string()]);
740
741 let out = p.execute(ctx, &HashMap::new()).await.unwrap();
742 assert_eq!(out.context.response.body.as_ref(), b"foo body");
744 assert_eq!(
745 out.context.response.headers.get("content-encoding"),
746 Some(&vec!["zstd".to_string()])
747 );
748 assert!(out.context.response.headers.contains_key("content-length"));
749 }
750
751 #[tokio::test]
752 async fn test_response_rewrite_filters_skip_corrupt_encoded_body() {
753 let p = plugin(serde_json::json!({
754 "filters": [{ "regex": "foo", "replace": "bar" }]
755 }));
756 let mut ctx = test_context(200, b"\x00not gzip\xff");
757 ctx.response
758 .headers
759 .insert("content-encoding".to_string(), vec!["gzip".to_string()]);
760
761 let out = p.execute(ctx, &HashMap::new()).await.unwrap();
762 assert_eq!(out.context.response.body.as_ref(), b"\x00not gzip\xff");
763 assert!(out
764 .context
765 .response
766 .headers
767 .contains_key("content-encoding"));
768 }
769
770 #[tokio::test]
771 async fn test_response_rewrite_vars_gate() {
772 let config = serde_json::json!({
773 "status_code": 500,
774 "body": "rewritten",
775 "headers": { "set": { "x-hit": "1" } },
776 "vars": [["status", "==", "200"]]
777 });
778
779 let p = plugin(config.clone());
781 let out = p
782 .execute(test_context(200, b"orig"), &HashMap::new())
783 .await
784 .unwrap();
785 assert_eq!(out.context.response.status_code, 500);
786 assert_eq!(out.context.response.body.as_ref(), b"rewritten");
787
788 let p = plugin(config);
790 let out = p
791 .execute(test_context(404, b"orig"), &HashMap::new())
792 .await
793 .unwrap();
794 assert_eq!(out.context.response.status_code, 404);
795 assert_eq!(out.context.response.body.as_ref(), b"orig");
796 assert!(!out.context.response.headers.contains_key("x-hit"));
797 assert!(out.context.response.headers.contains_key("content-length"));
798 }
799}