From 239236ccefc8e239f393b1327df660d48656789b Mon Sep 17 00:00:00 2001
From: favalos <gavalosf@gmail.com>
Date: Thu, 26 Sep 2024 10:17:39 -0700
Subject: [PATCH] Add a basic Url rewrite modifier.

---
 source/river/assets/test-config.kdl         |  1 +
 source/river/assets/test-config.toml        |  2 +
 source/river/src/proxy/mod.rs               |  3 ++
 source/river/src/proxy/request_modifiers.rs | 49 ++++++++++++++++++++-
 4 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/source/river/assets/test-config.kdl b/source/river/assets/test-config.kdl
index d6e1e2c..43cbdd4 100644
--- a/source/river/assets/test-config.kdl
+++ b/source/river/assets/test-config.kdl
@@ -96,6 +96,7 @@ services {
             upstream-request {
                 filter kind="remove-header-key-regex" pattern=".*(secret|SECRET).*"
                 filter kind="upsert-header" key="x-proxy-friend" value="river"
+                filter kind="url-rewrite" regex="^/rewrite/(.*)$" rewrite="/${1}"
             }
             upstream-response {
                 filter kind="remove-header-key-regex" pattern=".*ETag.*"
diff --git a/source/river/assets/test-config.toml b/source/river/assets/test-config.toml
index 37a6aff..b59109a 100644
--- a/source/river/assets/test-config.toml
+++ b/source/river/assets/test-config.toml
@@ -63,6 +63,8 @@ name = "Example1"
     # Filters are applied in the order they are specified. Multiple instances of
     # each filter may be provided.
     upstream-request-filters = [
+        # Url rewrite to
+        { kind = "url-rewrite", regex = "^/rewrite/(.*)$", rewrite = "/${1}" },
         # Remove any headers with keys matching `pattern`
         { kind = "remove-header-key-regex", pattern = ".*(secret|SECRET).*" },
         # Add or replace (e.g. "Upsert") a fixed header with the given key and value
diff --git a/source/river/src/proxy/mod.rs b/source/river/src/proxy/mod.rs
index b875d88..5491ed2 100644
--- a/source/river/src/proxy/mod.rs
+++ b/source/river/src/proxy/mod.rs
@@ -214,6 +214,9 @@ impl Modifiers {
                 "upsert-header" => {
                     Box::new(request_modifiers::UpsertHeader::from_settings(filter).unwrap())
                 }
+                "url-rewrite" => {
+                    Box::new(request_modifiers::PathRewrite::from_settings(filter).unwrap())
+                }
                 other => {
                     tracing::warn!("Unknown upstream request filter: '{other}'");
                     return Err(Error::new(ErrorType::Custom("Bad configuration")));
diff --git a/source/river/src/proxy/request_modifiers.rs b/source/river/src/proxy/request_modifiers.rs
index 1b4144f..3c37b4f 100644
--- a/source/river/src/proxy/request_modifiers.rs
+++ b/source/river/src/proxy/request_modifiers.rs
@@ -1,6 +1,8 @@
-use std::collections::BTreeMap;
+use std::{collections::BTreeMap, str::FromStr};
 
 use async_trait::async_trait;
+use http::{uri::PathAndQuery, Uri};
+use log::info;
 use pingora_core::{Error, Result};
 use pingora_http::RequestHeader;
 use pingora_proxy::Session;
@@ -112,3 +114,48 @@ impl RequestModifyMod for UpsertHeader {
         Ok(())
     }
 }
+
+pub struct PathRewrite {
+    regex: Regex,
+    rewrite: String,
+}
+
+impl PathRewrite {
+    // Create from settings
+    pub fn from_settings(mut settings: BTreeMap<String, String>) -> Result<Self> {
+        let regex = extract_val("regex", &mut settings)?;
+        let regex = Regex::from_str(regex.as_str()).expect("Unable to parse regex for rewrite.");
+        let rewrite = extract_val("rewrite", &mut settings)?;
+
+        Ok(Self {
+            regex,
+            rewrite: rewrite.clone(),
+        })
+    }
+}
+
+#[async_trait]
+impl RequestModifyMod for PathRewrite {
+    async fn upstream_request_filter(
+        &self,
+        _session: &mut Session,
+        header: &mut RequestHeader,
+        _ctx: &mut RiverContext,
+    ) -> Result<()> {
+        let path = header.uri.path();
+        match self.regex.is_match(path) {
+            false => Ok(()),
+            true => {
+                let rewrite = self.regex.replace(path, self.rewrite.as_str());
+
+                let rewritten_uri = Uri::builder()
+                    .path_and_query(PathAndQuery::from_str(&rewrite.to_string()).unwrap())
+                    .build();
+
+                header.set_uri(rewritten_uri.unwrap());
+
+                Ok(())
+            }
+        }
+    }
+}