Another refactor

Author: aram356

crates/common/Cargo.toml

@@ -51,6 +51,7 @@ config = { workspace = true }
 derive_more = { workspace = true }
 error-stack = { workspace = true }
 http = { workspace = true }
+log = { workspace = true }
 regex = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }

crates/common/src/html_processor.rs

@@ -65,7 +65,7 @@ impl StreamProcessor for HtmlWithPostProcessing {
         }
 
         if changed {
-            log::info!(
+            log::debug!(
                 "HTML post-processing complete: origin_host={}, output_len={}",
                 self.origin_host,
                 html.len()
@@ -138,6 +138,37 @@ pub fn create_html_processor(config: HtmlProcessorConfig) -> impl StreamProcesso
         fn protocol_relative_replacement(&self) -> String {
             format!("//{}", self.request_host)
         }
+
+        fn rewrite_url_value(&self, value: &str) -> Option<String> {
+            if !value.contains(&self.origin_host) {
+                return None;
+            }
+
+            let https_origin = self.https_origin();
+            let http_origin = self.http_origin();
+            let protocol_relative_origin = self.protocol_relative_origin();
+            let replacement_url = self.replacement_url();
+            let protocol_relative_replacement = self.protocol_relative_replacement();
+
+            let mut rewritten = value
+                .replace(&https_origin, &replacement_url)
+                .replace(&http_origin, &replacement_url)
+                .replace(&protocol_relative_origin, &protocol_relative_replacement);
+
+            if rewritten.starts_with(&self.origin_host) {
+                let suffix = &rewritten[self.origin_host.len()..];
+                let boundary_ok = suffix.is_empty()
+                    || matches!(
+                        suffix.as_bytes().first(),
+                        Some(b'/') | Some(b'?') | Some(b'#')
+                    );
+                if boundary_ok {
+                    rewritten = format!("{}{}", self.request_host, suffix);
+                }
+            }
+
+            (rewritten != value).then_some(rewritten)
+        }
     }
 
     let patterns = Rc::new(UrlPatterns {
@@ -170,11 +201,8 @@ pub fn create_html_processor(config: HtmlProcessorConfig) -> impl StreamProcesso
             move |el| {
                 if let Some(mut href) = el.get_attribute("href") {
                     let original_href = href.clone();
-                    let new_href = href
-                        .replace(&patterns.https_origin(), &patterns.replacement_url())
-                        .replace(&patterns.http_origin(), &patterns.replacement_url());
-                    if new_href != href {
-                        href = new_href;
+                    if let Some(rewritten) = patterns.rewrite_url_value(&href) {
+                        href = rewritten;
                     }
 
                     match integrations.rewrite_attribute(
@@ -211,11 +239,8 @@ pub fn create_html_processor(config: HtmlProcessorConfig) -> impl StreamProcesso
             move |el| {
                 if let Some(mut src) = el.get_attribute("src") {
                     let original_src = src.clone();
-                    let new_src = src
-                        .replace(&patterns.https_origin(), &patterns.replacement_url())
-                        .replace(&patterns.http_origin(), &patterns.replacement_url());
-                    if new_src != src {
-                        src = new_src;
+                    if let Some(rewritten) = patterns.rewrite_url_value(&src) {
+                        src = rewritten;
                     }
                     match integrations.rewrite_attribute(
                         "src",
@@ -251,11 +276,8 @@ pub fn create_html_processor(config: HtmlProcessorConfig) -> impl StreamProcesso
             move |el| {
                 if let Some(mut action) = el.get_attribute("action") {
                     let original_action = action.clone();
-                    let new_action = action
-                        .replace(&patterns.https_origin(), &patterns.replacement_url())
-                        .replace(&patterns.http_origin(), &patterns.replacement_url());
-                    if new_action != action {
-                        action = new_action;
+                    if let Some(rewritten) = patterns.rewrite_url_value(&action) {
+                        action = rewritten;
                     }
 
                     match integrations.rewrite_attribute(
@@ -510,8 +532,12 @@ mod tests {
 
         let html = r#"<html>
             <a href="https://origin.example.com/page">Link</a>
+            <a href="//origin.example.com/proto">Proto</a>
+            <a href="origin.example.com/bare">Bare</a>
             <img src="http://origin.example.com/image.jpg">
+            <img src="//origin.example.com/image2.jpg">
             <form action="https://origin.example.com/submit">
+            <form action="//origin.example.com/submit2">
         </html>"#;
 
         let mut output = Vec::new();
@@ -521,8 +547,12 @@ mod tests {
 
         let result = String::from_utf8(output).unwrap();
         assert!(result.contains(r#"href="https://test.example.com/page""#));
+        assert!(result.contains(r#"href="//test.example.com/proto""#));
+        assert!(result.contains(r#"href="test.example.com/bare""#));
         assert!(result.contains(r#"src="https://test.example.com/image.jpg""#));
+        assert!(result.contains(r#"src="//test.example.com/image2.jpg""#));
         assert!(result.contains(r#"action="https://test.example.com/submit""#));
+        assert!(result.contains(r#"action="//test.example.com/submit2""#));
         assert!(!result.contains("origin.example.com"));
     }