Refactor integration for post processing

aram356 · aram356 · commit 1fde992b9621 · 2025-12-15T10:56:51.000-08:00
diff --git a/crates/common/src/html_processor.rs b/crates/common/src/html_processor.rs
@@ -2,18 +2,84 @@
 //!
 //! This module provides a StreamProcessor implementation for HTML content.
 use std::cell::Cell;
+use std::io;
 use std::rc::Rc;
+use std::sync::Arc;
 
 use lol_html::{element, html_content::ContentType, text, Settings as RewriterSettings};
 
 use crate::integrations::{
-    AttributeRewriteOutcome, IntegrationAttributeContext, IntegrationRegistry,
-    IntegrationScriptContext, ScriptRewriteAction,
+    AttributeRewriteOutcome, IntegrationAttributeContext, IntegrationHtmlContext,
+    IntegrationHtmlPostProcessor, IntegrationRegistry, IntegrationScriptContext,
+    ScriptRewriteAction,
 };
 use crate::settings::Settings;
 use crate::streaming_processor::{HtmlRewriterAdapter, StreamProcessor};
 use crate::tsjs;
 
+struct HtmlWithPostProcessing {
+    inner: HtmlRewriterAdapter,
+    post_processors: Vec<Arc<dyn IntegrationHtmlPostProcessor>>,
+    origin_host: String,
+    request_host: String,
+    request_scheme: String,
+}
+
+impl StreamProcessor for HtmlWithPostProcessing {
+    fn process_chunk(&mut self, chunk: &[u8], is_last: bool) -> Result<Vec<u8>, io::Error> {
+        let output = self.inner.process_chunk(chunk, is_last)?;
+        if !is_last || output.is_empty() || self.post_processors.is_empty() {
+            return Ok(output);
+        }
+
+        let Ok(output_str) = std::str::from_utf8(&output) else {
+            return Ok(output);
+        };
+
+        let ctx = IntegrationHtmlContext {
+            request_host: &self.request_host,
+            request_scheme: &self.request_scheme,
+            origin_host: &self.origin_host,
+        };
+
+        // Preflight to avoid allocating a `String` unless at least one post-processor wants to run.
+        if !self
+            .post_processors
+            .iter()
+            .any(|p| p.should_process(output_str, &ctx))
+        {
+            return Ok(output);
+        }
+
+        let mut html = String::from_utf8(output).map_err(|e| {
+            io::Error::other(format!(
+                "HTML post-processing expected valid UTF-8 output: {e}"
+            ))
+        })?;
+
+        let mut changed = false;
+        for processor in &self.post_processors {
+            if processor.should_process(&html, &ctx) {
+                changed |= processor.post_process(&mut html, &ctx);
+            }
+        }
+
+        if changed {
+            log::info!(
+                "HTML post-processing complete: origin_host={}, output_len={}",
+                self.origin_host,
+                html.len()
+            );
+        }
+
+        Ok(html.into_bytes())
+    }
+
+    fn reset(&mut self) {
+        self.inner.reset();
+    }
+}
+
 /// Configuration for HTML processing
 #[derive(Clone)]
 pub struct HtmlProcessorConfig {
@@ -43,6 +109,8 @@ impl HtmlProcessorConfig {
 
 /// Create an HTML processor with URL replacement and optional Prebid injection
 pub fn create_html_processor(config: HtmlProcessorConfig) -> impl StreamProcessor {
+    let post_processors = config.integrations.html_post_processors();
+
     // Simplified URL patterns structure - stores only core data and generates variants on-demand
     struct UrlPatterns {
         origin_host: String,
@@ -343,7 +411,13 @@ pub fn create_html_processor(config: HtmlProcessorConfig) -> impl StreamProcesso
         ..RewriterSettings::default()
     };
 
-    HtmlRewriterAdapter::new(rewriter_settings)
+    HtmlWithPostProcessing {
+        inner: HtmlRewriterAdapter::new(rewriter_settings),
+        post_processors,
+        origin_host: config.origin_host,
+        request_host: config.request_host,
+        request_scheme: config.request_scheme,
+    }
 }
 
 #[cfg(test)]
diff --git a/crates/common/src/integrations/nextjs.rs b/crates/common/src/integrations/nextjs.rs
@@ -95,43 +95,40 @@ impl IntegrationHtmlPostProcessor for NextJsHtmlPostProcessor {
         NEXTJS_INTEGRATION_ID
     }
 
-    fn post_process(&self, html: &str, ctx: &IntegrationHtmlContext<'_>) -> String {
-        log::info!(
-            "NextJs post-processor called: enabled={}, rewrite_attributes={:?}, html_len={}, origin={}, proxy={}://{}",
-            self.config.enabled,
-            self.config.rewrite_attributes,
-            html.len(),
-            ctx.origin_host,
-            ctx.request_scheme,
-            ctx.request_host
-        );
-
+    fn should_process(&self, html: &str, ctx: &IntegrationHtmlContext<'_>) -> bool {
         if !self.config.enabled || self.config.rewrite_attributes.is_empty() {
-            log::info!("NextJs post-processor skipped (disabled or no attributes)");
-            return html.to_string();
+            return false;
+        }
+
+        // Only Next.js App Router pages will contain `__next_f` pushes.
+        // Also require an origin host hit to avoid running on already-rewritten pages.
+        html.contains("__next_f.push") && html.contains(ctx.origin_host)
+    }
+
+    fn post_process(&self, html: &mut String, ctx: &IntegrationHtmlContext<'_>) -> bool {
+        if !self.should_process(html, ctx) {
+            return false;
         }
 
-        // Count origin URLs before
         let origin_before = html.matches(ctx.origin_host).count();
         log::info!(
-            "NextJs post-processor: {} origin URLs before rewrite",
-            origin_before
+            "NextJs post-processor running: html_len={}, origin_matches={}, origin={}, proxy={}://{}",
+            html.len(),
+            origin_before,
+            ctx.origin_host,
+            ctx.request_scheme,
+            ctx.request_host
         );
 
         let result =
             post_process_rsc_html(html, ctx.origin_host, ctx.request_host, ctx.request_scheme);
 
-        // Count after
-        let origin_after = result.matches(ctx.origin_host).count();
-        let proxy_after = result.matches(ctx.request_host).count();
-        log::info!(
-            "NextJs post-processor complete: input_len={}, output_len={}, origin_remaining={}, proxy_urls={}",
-            html.len(),
-            result.len(),
-            origin_after,
-            proxy_after
-        );
-        result
+        if result == *html {
+            return false;
+        }
+
+        *html = result;
+        true
     }
 }
 
@@ -291,14 +288,14 @@ impl IntegrationScriptRewriter for NextJsScriptRewriter {
         match self.mode {
             NextJsRewriteMode::Structured => self.rewrite_structured(content, ctx),
             NextJsRewriteMode::Streamed => {
-                // RSC push scripts (self.__next_f.push) are handled by the post-processor
+                // RSC push scripts (__next_f.push) are handled by the post-processor
                 // because T-chunks can span multiple scripts and require combined processing.
                 // Only handle non-RSC scripts here.
-                if content.contains("self.__next_f.push") {
+                if content.contains("__next_f.push") {
                     return ScriptRewriteAction::keep();
                 }
                 // For other __next_f scripts (like initialization), use simple URL rewriting
-                if content.contains("self.__next_f") {
+                if content.contains("__next_f") {
                     return self.rewrite_streamed(content, ctx);
                 }
                 ScriptRewriteAction::keep()
diff --git a/crates/common/src/integrations/registry.rs b/crates/common/src/integrations/registry.rs
@@ -264,10 +264,19 @@ pub trait IntegrationHtmlPostProcessor: Send + Sync {
     /// Identifier for logging/diagnostics.
     fn integration_id(&self) -> &'static str;
 
+    /// Fast preflight check to decide whether post-processing should run for this document.
+    ///
+    /// Implementations should keep this cheap (e.g., a substring check) because it may run on
+    /// every HTML response when the integration is enabled.
+    fn should_process(&self, html: &str, ctx: &IntegrationHtmlContext<'_>) -> bool {
+        let _ = (html, ctx);
+        true
+    }
+
     /// Post-process complete HTML content.
     /// This is called after streaming HTML processing with the complete HTML.
-    /// Return the modified HTML or the original if no changes needed.
-    fn post_process(&self, html: &str, ctx: &IntegrationHtmlContext<'_>) -> String;
+    /// Implementations should mutate `html` in-place and return `true` when changes were made.
+    fn post_process(&self, html: &mut String, ctx: &IntegrationHtmlContext<'_>) -> bool;
 }
 
 /// Registration payload returned by integration builders.
diff --git a/crates/common/src/publisher.rs b/crates/common/src/publisher.rs
@@ -1,75 +1,20 @@
 use error_stack::{Report, ResultExt};
 use fastly::http::{header, StatusCode};
 use fastly::{Body, Request, Response};
-use std::io::Write;
 
 use crate::backend::ensure_backend_from_url;
 use crate::http_util::serve_static_with_etag;
 
 use crate::constants::{HEADER_SYNTHETIC_TRUSTED_SERVER, HEADER_X_COMPRESS_HINT};
 use crate::cookies::create_synthetic_cookie;
 use crate::error::TrustedServerError;
-use crate::integrations::{IntegrationHtmlContext, IntegrationRegistry};
+use crate::integrations::IntegrationRegistry;
 use crate::rsc_flight::RscFlightUrlRewriter;
 use crate::settings::Settings;
 use crate::streaming_processor::{Compression, PipelineConfig, StreamProcessor, StreamingPipeline};
 use crate::streaming_replacer::create_url_replacer;
 use crate::synthetic::get_or_generate_synthetic_id;
 
-/// Compress data using the specified compression algorithm
-fn compress_data(
-    data: &[u8],
-    compression: Compression,
-) -> Result<Vec<u8>, Report<TrustedServerError>> {
-    match compression {
-        Compression::None => Ok(data.to_vec()),
-        Compression::Gzip => {
-            use flate2::write::GzEncoder;
-            use flate2::Compression as GzCompression;
-            let mut encoder = GzEncoder::new(Vec::new(), GzCompression::default());
-            encoder
-                .write_all(data)
-                .change_context(TrustedServerError::Proxy {
-                    message: "Failed to gzip compress data".to_string(),
-                })?;
-            encoder.finish().change_context(TrustedServerError::Proxy {
-                message: "Failed to finish gzip compression".to_string(),
-            })
-        }
-        Compression::Deflate => {
-            use flate2::write::ZlibEncoder;
-            use flate2::Compression as ZlibCompression;
-            let mut encoder = ZlibEncoder::new(Vec::new(), ZlibCompression::default());
-            encoder
-                .write_all(data)
-                .change_context(TrustedServerError::Proxy {
-                    message: "Failed to deflate compress data".to_string(),
-                })?;
-            encoder.finish().change_context(TrustedServerError::Proxy {
-                message: "Failed to finish deflate compression".to_string(),
-            })
-        }
-        Compression::Brotli => {
-            use brotli::enc::writer::CompressorWriter;
-            use brotli::enc::BrotliEncoderParams;
-            let params = BrotliEncoderParams {
-                quality: 4, // Balance speed and compression
-                ..Default::default()
-            };
-            let mut output = Vec::new();
-            {
-                let mut writer = CompressorWriter::with_params(&mut output, 4096, &params);
-                writer
-                    .write_all(data)
-                    .change_context(TrustedServerError::Proxy {
-                        message: "Failed to brotli compress data".to_string(),
-                    })?;
-            }
-            Ok(output)
-        }
-    }
-}
-
 /// Detects the request scheme (HTTP or HTTPS) using Fastly SDK methods and headers.
 ///
 /// Tries multiple methods in order of reliability:
@@ -199,67 +144,14 @@ fn process_response_streaming(
             params.integration_registry,
         )?;
 
-        // Check if we have post-processors that need uncompressed HTML
-        let post_processors = params.integration_registry.html_post_processors();
-        let needs_post_processing = !post_processors.is_empty();
-
-        // If we have post-processors, output uncompressed HTML so they can work with it,
-        // then compress only once at the end. This avoids double decompression/compression.
-        let output_compression = if needs_post_processing {
-            Compression::None
-        } else {
-            compression
-        };
-
         let config = PipelineConfig {
             input_compression: compression,
-            output_compression,
+            output_compression: compression,
             chunk_size: 8192,
         };
 
         let mut pipeline = StreamingPipeline::new(config, processor);
         pipeline.process(body, &mut output)?;
-
-        // Post-process HTML through registered integration post-processors.
-        // This handles cross-script T-chunks for RSC and other integration-specific
-        // processing that requires the complete HTML document.
-        log::info!(
-            "HTML post-processors: count={}, output_len={}, needs_post_processing={}",
-            post_processors.len(),
-            output.len(),
-            needs_post_processing
-        );
-        if needs_post_processing {
-            // Output is already uncompressed, convert to string for post-processing
-            if let Ok(html) = std::str::from_utf8(&output) {
-                log::info!(
-                    "NextJs post-processor called with {} bytes of HTML",
-                    html.len()
-                );
-                let ctx = IntegrationHtmlContext {
-                    request_host: params.request_host,
-                    request_scheme: params.request_scheme,
-                    origin_host: params.origin_host,
-                };
-                let mut processed = html.to_string();
-                for processor in post_processors {
-                    processed = processor.post_process(&processed, &ctx);
-                }
-
-                // Now compress if original content was compressed
-                if compression != Compression::None {
-                    output = compress_data(processed.as_bytes(), compression)?;
-                } else {
-                    output = processed.into_bytes();
-                }
-            } else {
-                log::warn!("HTML post-processing skipped: content is not valid UTF-8");
-                // If not valid UTF-8, recompress the output as-is
-                if compression != Compression::None {
-                    output = compress_data(&output, compression)?;
-                }
-            }
-        }
     } else if is_rsc_flight {
         // RSC Flight responses are length-prefixed (T rows). A naive string replacement will
         // corrupt the stream by changing byte lengths without updating the prefixes.
diff --git a/docs/RSC_HYDRATION_FINDINGS.md b/docs/RSC_HYDRATION_FINDINGS.md
diff --git a/docs/integration_guide.md b/docs/integration_guide.md
