Another refactor

aram356 · aram356 · commit 86a807b1e27d · 2025-12-15T21:30:41.000-08:00
diff --git a/crates/common/src/integrations/nextjs/html_post_process.rs b/crates/common/src/integrations/nextjs/html_post_process.rs
@@ -84,16 +84,19 @@ fn find_rsc_push_scripts(html: &str) -> Vec<RscPushScriptRange> {
         let mut i = payload_start;
         let bytes = html.as_bytes();
         while i < bytes.len() {
-            if bytes[i] == b'\\' {
-                i += 2;
+            if bytes[i] == b'\\' && i + 1 < bytes.len() {
+                i += 2; // Skip escape sequence (safe: we checked i+1 exists)
+            } else if bytes[i] == b'\\' {
+                // Trailing backslash at end of content - malformed
+                break;
             } else if bytes[i] == quote as u8 {
                 break;
             } else {
                 i += 1;
             }
         }
 
-        if i >= bytes.len() {
+        if i >= bytes.len() || bytes[i] != quote as u8 {
             search_pos = payload_start;
             continue;
         }
@@ -330,4 +333,54 @@ mod tests {
             "Origin URL should be removed. Got: {rewritten}"
         );
     }
+
+    #[test]
+    fn handles_trailing_backslash_gracefully() {
+        // Malformed content with trailing backslash should not panic
+        let html = r#"<html><body>
+<script>self.__next_f.push([1,"content with trailing backslash\"])</script>
+<script>self.__next_f.push([1,"valid https://origin.example.com/page"])</script>
+</body></html>"#;
+
+        let scripts = find_rsc_push_scripts(html);
+        // The first script is malformed (trailing backslash escapes the quote),
+        // so it won't be detected as valid. The second one should be found.
+        assert!(
+            scripts.len() >= 1,
+            "Should find at least the valid script. Found: {}",
+            scripts.len()
+        );
+
+        // Should not panic during processing
+        let result = post_process_rsc_html(html, "origin.example.com", "test.example.com", "https");
+        assert!(
+            result.contains("test.example.com") || result.contains("origin.example.com"),
+            "Processing should complete without panic"
+        );
+    }
+
+    #[test]
+    fn handles_unterminated_string_gracefully() {
+        // Content where string never closes - should not hang or panic
+        let html = r#"<html><body>
+<script>self.__next_f.push([1,"content without closing quote
+</body></html>"#;
+
+        let scripts = find_rsc_push_scripts(html);
+        assert_eq!(
+            scripts.len(),
+            0,
+            "Should not find scripts with unterminated strings"
+        );
+    }
+
+    #[test]
+    fn no_origin_returns_unchanged() {
+        let html = r#"<html><body>
+<script>self.__next_f.push([1,"content without origin URLs"])</script>
+</body></html>"#;
+
+        let result = post_process_rsc_html(html, "origin.example.com", "test.example.com", "https");
+        assert_eq!(result, html, "HTML without origin should be unchanged");
+    }
 }
diff --git a/crates/common/src/integrations/nextjs/rsc.rs b/crates/common/src/integrations/nextjs/rsc.rs
@@ -10,6 +10,10 @@ static TCHUNK_PATTERN: Lazy<Regex> =
 /// Marker used to track script boundaries when combining RSC content.
 pub(crate) const RSC_MARKER: &str = "\x00SPLIT\x00";
 
+/// Maximum combined payload size for cross-script processing (10 MB).
+/// Payloads exceeding this limit are processed individually without cross-script T-chunk handling.
+const MAX_COMBINED_PAYLOAD_SIZE: usize = 10 * 1024 * 1024;
+
 // =============================================================================
 // Escape Sequence Parsing
 // =============================================================================
@@ -299,6 +303,7 @@ impl RscUrlRewriter {
             return Cow::Borrowed(input);
         }
 
+        // Phase 1: Regex-based URL pattern rewriting (handles escaped slashes, schemes, etc.)
         let replaced = self
             .pattern
             .replace_all(input, |caps: &regex::Captures<'_>| {
@@ -310,18 +315,20 @@ impl RscUrlRewriter {
                 }
             });
 
-        let still_contains_origin = match &replaced {
-            Cow::Borrowed(s) => s.contains(&self.origin_host),
-            Cow::Owned(s) => s.contains(&self.origin_host),
+        // Phase 2: Handle bare host occurrences not matched by the URL regex
+        // (e.g., `siteProductionDomain`). Only check if regex made no changes,
+        // because if it did, we already know origin_host was present.
+        let text = match &replaced {
+            Cow::Borrowed(s) => *s,
+            Cow::Owned(s) => s.as_str(),
         };
 
-        if !still_contains_origin {
+        if !text.contains(&self.origin_host) {
             return replaced;
         }
 
-        // Also rewrite bare host occurrences inside RSC payloads (e.g. `siteProductionDomain`).
-        let owned = replaced.into_owned();
-        Cow::Owned(owned.replace(&self.origin_host, &self.request_host))
+        // Bare host replacement needed
+        Cow::Owned(text.replace(&self.origin_host, &self.request_host))
     }
 
     pub(crate) fn rewrite_to_string(&self, input: &str) -> String {
@@ -398,7 +405,26 @@ pub fn rewrite_rsc_scripts_combined(
         return vec![rewrite_rsc_tchunks_with_rewriter(payloads[0], &rewriter)];
     }
 
-    let mut combined = payloads[0].to_string();
+    // Check total size before allocating combined buffer
+    let total_size: usize =
+        payloads.iter().map(|p| p.len()).sum::<usize>() + (payloads.len() - 1) * RSC_MARKER.len();
+
+    if total_size > MAX_COMBINED_PAYLOAD_SIZE {
+        // Fall back to individual processing if combined size is too large.
+        // This sacrifices cross-script T-chunk correctness for memory safety.
+        log::warn!(
+            "RSC combined payload size {} exceeds limit {}, processing individually",
+            total_size,
+            MAX_COMBINED_PAYLOAD_SIZE
+        );
+        return payloads
+            .iter()
+            .map(|p| rewrite_rsc_tchunks_with_rewriter(p, &rewriter))
+            .collect();
+    }
+
+    let mut combined = String::with_capacity(total_size);
+    combined.push_str(payloads[0]);
     for payload in &payloads[1..] {
         combined.push_str(RSC_MARKER);
         combined.push_str(payload);
@@ -591,4 +617,80 @@ mod tests {
             "Bare host should be rewritten inside RSC payload. Got: {rewritten}"
         );
     }
+
+    #[test]
+    fn single_payload_bypasses_combining() {
+        // When there's only one payload, we should process it directly without combining
+        // Content: {"url":"https://origin.example.com/x"} = 37 bytes = 0x25 hex
+        let payload = r#"1a:T25,{"url":"https://origin.example.com/x"}"#;
+        let payloads: Vec<&str> = vec![payload];
+
+        let results = rewrite_rsc_scripts_combined(
+            &payloads,
+            "origin.example.com",
+            "test.example.com",
+            "https",
+        );
+
+        assert_eq!(results.len(), 1);
+        assert!(
+            results[0].contains("test.example.com"),
+            "Single payload should be rewritten. Got: {}",
+            results[0]
+        );
+        // The length should be updated for the rewritten URL
+        // {"url":"https://test.example.com/x"} = 35 bytes = 0x23 hex
+        assert!(
+            results[0].contains(":T23,"),
+            "T-chunk length should be updated. Got: {}",
+            results[0]
+        );
+    }
+
+    #[test]
+    fn empty_payloads_returns_empty() {
+        let payloads: Vec<&str> = vec![];
+        let results = rewrite_rsc_scripts_combined(
+            &payloads,
+            "origin.example.com",
+            "test.example.com",
+            "https",
+        );
+        assert!(results.is_empty());
+    }
+
+    #[test]
+    fn no_origin_in_payloads_returns_unchanged() {
+        let payloads: Vec<&str> = vec![r#"1a:T10,{"key":"value"}"#, r#"1b:T10,{"foo":"bar"}"#];
+
+        let results = rewrite_rsc_scripts_combined(
+            &payloads,
+            "origin.example.com",
+            "test.example.com",
+            "https",
+        );
+
+        assert_eq!(results.len(), 2);
+        // Content should be identical - note that T-chunk lengths may be recalculated
+        // even if content is unchanged (due to how the algorithm works)
+        assert!(
+            !results[0].contains("origin.example.com") && !results[0].contains("test.example.com"),
+            "No host should be present in payload without URLs"
+        );
+        assert!(
+            !results[1].contains("origin.example.com") && !results[1].contains("test.example.com"),
+            "No host should be present in payload without URLs"
+        );
+        // The content after T-chunk header should be preserved
+        assert!(
+            results[0].contains(r#"{"key":"value"}"#),
+            "Content should be preserved. Got: {}",
+            results[0]
+        );
+        assert!(
+            results[1].contains(r#"{"foo":"bar"}"#),
+            "Content should be preserved. Got: {}",
+            results[1]
+        );
+    }
 }
diff --git a/crates/common/src/integrations/nextjs/script_rewriter.rs b/crates/common/src/integrations/nextjs/script_rewriter.rs
@@ -36,14 +36,22 @@ impl NextJsScriptRewriter {
         content: &str,
         ctx: &IntegrationScriptContext<'_>,
     ) -> ScriptRewriteAction {
-        if let Some(rewritten) = rewrite_nextjs_values(
-            content,
+        if ctx.origin_host.is_empty()
+            || ctx.request_host.is_empty()
+            || self.config.rewrite_attributes.is_empty()
+        {
+            return ScriptRewriteAction::keep();
+        }
+
+        let rewriter = UrlRewriter::new(
             ctx.origin_host,
             ctx.request_host,
             ctx.request_scheme,
             &self.config.rewrite_attributes,
-            false,
-        ) {
+            false, // preserve_length not used for structured payloads
+        );
+
+        if let Some(rewritten) = rewrite_nextjs_values_with_rewriter(content, &rewriter) {
             ScriptRewriteAction::replace(rewritten)
         } else {
             ScriptRewriteAction::keep()
@@ -151,6 +159,11 @@ impl IntegrationScriptRewriter for NextJsScriptRewriter {
     }
 }
 
+fn rewrite_nextjs_values_with_rewriter(content: &str, rewriter: &UrlRewriter) -> Option<String> {
+    rewriter.rewrite_embedded(content)
+}
+
+#[cfg(test)]
 fn rewrite_nextjs_values(
     content: &str,
     origin_host: &str,
@@ -171,15 +184,26 @@ fn rewrite_nextjs_values(
         preserve_length,
     );
 
-    rewriter.rewrite_embedded(content)
+    rewrite_nextjs_values_with_rewriter(content, &rewriter)
 }
 
+/// Rewrites URLs in structured Next.js JSON payloads (e.g., `__NEXT_DATA__`).
+///
+/// This rewriter uses attribute-specific regex patterns to find and replace URLs
+/// in JSON content. It handles full URLs, protocol-relative URLs, and bare hostnames.
+///
+/// The `preserve_length` option adds whitespace padding to maintain byte length,
+/// which was an early attempt at RSC compatibility. This is no longer needed for
+/// RSC payloads (T-chunk lengths are recalculated instead), but is kept for
+/// potential future use cases where length preservation is required.
 struct UrlRewriter {
     origin_host: String,
     request_host: String,
     request_scheme: String,
     embedded_patterns: Vec<Regex>,
     bare_host_patterns: Vec<Regex>,
+    /// When true, adds whitespace padding to maintain original byte length.
+    /// Currently unused in production (always false).
     preserve_length: bool,
 }
 
diff --git a/crates/common/src/rsc_flight.rs b/crates/common/src/rsc_flight.rs
@@ -21,6 +21,14 @@ enum RowState {
 ///
 /// For `T` rows, the length prefix is the UTF-8 byte length of the content bytes. If we rewrite
 /// URLs inside the content, we must recompute the length and rewrite the header.
+///
+/// ## Limitations
+///
+/// This rewriter performs simple string replacement and does NOT handle JSON escape sequences.
+/// URLs like `\/\/origin.example.com` (JSON-escaped slashes) will not be rewritten. This is
+/// acceptable because Flight responses from client-side navigation typically contain plain URLs,
+/// not doubly-escaped JSON-in-JS content. For inlined `__next_f` data in HTML (which can have
+/// escape sequences), the HTML post-processor in `integrations/nextjs/` handles those cases.
 pub struct RscFlightUrlRewriter {
     origin_url: String,
     origin_http_url: Option<String>,
diff --git a/docs/RSC_HYDRATION_FINDINGS.md b/docs/RSC_HYDRATION_FINDINGS.md
