Allow pre-parsing of ESI document, add example of fetching per-user fragment variants

Cargo.toml

@@ -5,6 +5,7 @@ members = [
     "examples/esi_example_minimal",
     "examples/esi_example_advanced_error_handling",
     "examples/esi_try_example",
+    "examples/esi_example_variants",
 ]
 
 [workspace.package]

README.md

@@ -52,7 +52,7 @@ fn handle_request(req: Request) -> Result<(), Error> {
             // of the request URL will be used as the backend name.
             Some(&|req| {
                 println!("Sending request {} {}", req.get_method(), req.get_path());
-                Ok(Some(req.with_ttl(120).send_async("mock-s3")?))
+                Ok(req.with_ttl(120).send_async("mock-s3")?.into())
             }),
             // Optionally provide a method to process fragment responses before they
             // are streamed to the client.

esi/src/lib.rs

@@ -111,6 +111,45 @@ impl Processor {
         }
     }
 
+    /// Process an ESI document that has already been parsed into a queue of events.
+    pub fn process_parsed_document(
+        self,
+        src_events: VecDeque<Event>,
+        output_writer: &mut Writer<impl Write>,
+        dispatch_fragment_request: Option<&FragmentRequestDispatcher>,
+        process_fragment_response: Option<&FragmentResponseProcessor>,
+    ) -> Result<()> {
+        // Set up fragment request dispatcher. Use what's provided or use a default
+        let dispatch_fragment_request =
+            dispatch_fragment_request.unwrap_or(&default_fragment_dispatcher);
+
+        // If there is a source request to mimic, copy its metadata, otherwise use a default request.
+        let original_request_metadata = self.original_request_metadata.as_ref().map_or_else(
+            || Request::new(Method::GET, "http://localhost"),
+            Request::clone_without_body,
+        );
+
+        // `root_task` is the root task that will be used to fetch tags in recursive manner
+        let root_task: &mut Task = &mut Task::new();
+
+        for event in src_events {
+            event_receiver(
+                event,
+                &mut root_task.queue,
+                self.configuration.is_escaped_content,
+                &original_request_metadata,
+                dispatch_fragment_request,
+            )?;
+        }
+
+        self.process_root_task(
+            root_task,
+            output_writer,
+            dispatch_fragment_request,
+            process_fragment_response,
+        )
+    }
+
     /// Process an ESI document from a [`quick_xml::Reader`].
     pub fn process_document(
         self,
@@ -120,18 +159,8 @@ impl Processor {
         process_fragment_response: Option<&FragmentResponseProcessor>,
     ) -> Result<()> {
         // Set up fragment request dispatcher. Use what's provided or use a default
-        let dispatch_fragment_request = dispatch_fragment_request.unwrap_or({
-            &|req| {
-                debug!("no dispatch method configured, defaulting to hostname");
-                let backend = req
-                    .get_url()
-                    .host()
-                    .unwrap_or_else(|| panic!("no host in request: {}", req.get_url()))
-                    .to_string();
-                let pending_req = req.send_async(backend)?;
-                Ok(pending_req.into())
-            }
-        });
+        let dispatch_fragment_request =
+            dispatch_fragment_request.unwrap_or(&default_fragment_dispatcher);
 
         // If there is a source request to mimic, copy its metadata, otherwise use a default request.
         let original_request_metadata = self.original_request_metadata.as_ref().map_or_else(
@@ -140,9 +169,8 @@ impl Processor {
         );
 
         // `root_task` is the root task that will be used to fetch tags in recursive manner
-        let root_task = &mut Task::new();
+        let root_task: &mut Task = &mut Task::new();
 
-        let is_escaped = self.configuration.is_escaped_content;
         // Call the library to parse fn `parse_tags` which will call the callback function
         // on each tag / event it finds in the document.
         // The callback function `handle_events` will handle the event.
@@ -153,18 +181,34 @@ impl Processor {
                 event_receiver(
                     event,
                     &mut root_task.queue,
-                    is_escaped,
+                    self.configuration.is_escaped_content,
                     &original_request_metadata,
                     dispatch_fragment_request,
                 )
             },
         )?;
 
+        self.process_root_task(
+            root_task,
+            output_writer,
+            dispatch_fragment_request,
+            process_fragment_response,
+        )
+    }
+
+    fn process_root_task(
+        self,
+        root_task: &mut Task,
+        output_writer: &mut Writer<impl Write>,
+        dispatch_fragment_request: &FragmentRequestDispatcher,
+        process_fragment_response: Option<&FragmentResponseProcessor>,
+    ) -> Result<()> {
         // set the root depth to 0
         let mut depth = 0;
 
         debug!("Elements to fetch: {:?}", root_task.queue);
-        // Elements dependent on backend requests got are queued up.
+
+        // Elements dependent on backend requests are queued up.
         // The responses will need to be fetched and processed.
         // Go over the list for any pending responses and write them to the client output stream.
         fetch_elements(
@@ -179,6 +223,17 @@ impl Processor {
     }
 }
 
+fn default_fragment_dispatcher(req: Request) -> Result<PendingFragmentContent> {
+    debug!("no dispatch method configured, defaulting to hostname");
+    let backend = req
+        .get_url()
+        .host()
+        .unwrap_or_else(|| panic!("no host in request: {}", req.get_url()))
+        .to_string();
+    let pending_req = req.send_async(backend)?;
+    Ok(PendingFragmentContent::PendingRequest(pending_req))
+}
+
 // This function is responsible for fetching pending requests and writing their
 // responses to the client output stream. It also handles any queued source
 // content that needs to be written to the client output stream.

examples/esi_example_variants/.cargo/config.toml

@@ -0,0 +1,6 @@
+[target.wasm32-wasi]
+rustflags = ["-C", "debuginfo=2"]
+runner = "viceroy run -C fastly.toml -- "
+
+[build]
+target = "wasm32-wasi"

examples/esi_example_variants/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "esi_example_variants"
+version.workspace = true
+authors.workspace = true
+license.workspace = true
+edition.workspace = true
+publish = false
+
+[dependencies]
+fastly = "^0.11"
+esi = { path = "../../esi" }
+log = "^0.4"
+env_logger = "^0.11"

examples/esi_example_variants/fastly.toml

@@ -0,0 +1,28 @@
+# This file describes a Fastly Compute package. To learn more visit:
+# https://developer.fastly.com/reference/fastly-toml/
+
+authors = ["kailan@enviark.com"]
+description = ""
+language = "rust"
+manifest_version = 2
+name = "esi_example_variants"
+service_id = ""
+
+[local_server]
+
+  [local_server.backends]
+
+    [local_server.backends.mock-s3]
+      url = "https://mock-s3.edgecompute.app"
+      override_host = "mock-s3.edgecompute.app"
+
+[scripts]
+  build = "cargo build --bin esi_example_variants --release --target wasm32-wasi --color always"
+
+[setup]
+
+  [setup.backends]
+
+    [setup.backends.mock-s3]
+      address = "mock-s3.edgecompute.app"
+      port = 443

examples/esi_example_variants/src/index.html

@@ -0,0 +1,15 @@
+<!DOCTYPE html>
+<html>
+  <head>
+    <title>My Broken Website</title>
+  </head>
+  <body>
+    <header style="background: #f1f1f1; padding: 16px">
+      <h1>My Broken Website</h1>
+    </header>
+    <p>There is <em>some</em> valid content on this page, like the text you're reading right now.</p>
+    <p>But watch this... I'm about to include an ESI fragment that does not exist:</p>
+    <esi:include src="/_fragments/doesnotexist.html"/>
+    <p>If you're seeing this text in the browser, then something is broken.</p>
+  </body>
+</html>

examples/esi_example_variants/src/main.rs

@@ -0,0 +1,123 @@
+use std::{
+    collections::{HashMap, VecDeque},
+    io::Write,
+};
+
+use esi::{parse_tags, Reader, Writer};
+use fastly::{http::StatusCode, mime, Request, Response};
+use log::{error, info};
+
+// Take a list of fragment URLs and return a map that maps each URL to a variant for the specific request.
+fn get_variant_urls(urls: Vec<String>) -> HashMap<String, String> {
+    let mut variant_urls = HashMap::new();
+    for url in urls {
+        // For demonstration, add a query parameter to each request
+        let variant_url = if url.contains('?') {
+            format!("{}&variant=1", url)
+        } else {
+            format!("{}?variant=1", url)
+        };
+        variant_urls.insert(url, variant_url);
+    }
+    variant_urls
+}
+
+fn main() {
+    env_logger::builder()
+        .filter(None, log::LevelFilter::Trace)
+        .init();
+
+    let req = Request::from_client();
+
+    if req.get_path() != "/" {
+        Response::from_status(StatusCode::NOT_FOUND).send_to_client();
+        return;
+    }
+
+    // Generate synthetic test response from "index.html" file.
+    // You probably want replace this with a backend call, e.g. `req.clone_without_body().send("origin_0")`
+    let mut beresp =
+        Response::from_body(include_str!("index.html")).with_content_type(mime::TEXT_HTML);
+
+    // If the response is HTML, we can parse it for ESI tags.
+    if beresp
+        .get_content_type()
+        .is_some_and(|c| c.subtype() == mime::HTML)
+    {
+        let processor = esi::Processor::new(Some(req), esi::Configuration::default());
+
+        // Create a response to send the headers to the client
+        let resp = Response::from_status(StatusCode::OK).with_content_type(mime::TEXT_HTML);
+
+        // Send the response headers to the client and open an output stream
+        let output_writer = resp.stream_to_client();
+
+        // Set up an XML writer to write directly to the client output stream.
+        let mut xml_writer = Writer::new(output_writer);
+
+        // Parse the ESI document and store it in memory
+        let mut events = VecDeque::new();
+        let mut beresp_reader = Reader::from_reader(beresp.take_body());
+        parse_tags("esi", &mut beresp_reader, &mut |event| {
+            events.push_back(event);
+            Ok(())
+        })
+        .expect("failed to parse ESI template");
+
+        // Extract the `src` URLs from ESI includes
+        let urls = events
+            .iter()
+            .filter_map(|event| match event {
+                esi::Event::ESI(esi::Tag::Include { src, .. }) => Some(src.clone()),
+                _ => None,
+            })
+            .collect::<Vec<_>>();
+
+        // Check the variant database to determine the URLs to fetch
+        let variant_urls = get_variant_urls(urls);
+
+        // Process the already-parsed ESI document, replacing the request URLs with the variant URLs
+        let result = processor.process_parsed_document(
+            events,
+            &mut xml_writer,
+            Some(&move |req| {
+                let original_url = req.get_url().to_string();
+                let variant_url = variant_urls.get(&original_url).unwrap_or(&original_url);
+                info!(
+                    "Sending request - original URL: ({}) variant URL: ({})",
+                    req.get_path(),
+                    variant_url
+                );
+                Ok(esi::PendingFragmentContent::PendingRequest(
+                    req.with_url(variant_url)
+                        .with_ttl(120)
+                        .send_async("mock-s3")?,
+                ))
+            }),
+            Some(&|req, resp| {
+                info!(
+                    "Received response for {} {}",
+                    req.get_method(),
+                    req.get_path()
+                );
+                Ok(resp)
+            }),
+        );
+
+        match result {
+            Ok(()) => {
+                xml_writer.into_inner().finish().unwrap();
+            }
+            Err(err) => {
+                error!("error processing ESI document: {}", err);
+                let _ = xml_writer.get_mut().write(b"Internal server error");
+                xml_writer.into_inner().finish().unwrap_or_else(|_| {
+                    error!("error flushing error response to client");
+                });
+            }
+        }
+    } else {
+        // Otherwise, we can just return the response.
+        beresp.send_to_client();
+    }
+}

examples/esi_try_example/Cargo.toml

@@ -10,7 +10,7 @@ publish = false
 debug = 1
 
 [dependencies]
-fastly = "0.11"
+fastly = "^0.11"
 esi = { path = "../../esi" }
 log = "^0.4"
 env_logger = "0.11.3"