feat: add postBody Functionality

README.md

@@ -19,6 +19,7 @@ Library supports all instructions that the Ledge parser supports. Also supports
     - [varsCookieBlacklist](#varscookieblacklist)
   - [Custom ESI Vars Function](#custom-esi-vars-function)
   - [Custom Fetch Function](#custom-fetch-function)
+  - [Optional Post Body Function](#optional-post-body-function)
 - [Caching and upstream requests](#caching-and-upstream-requests)
 - [Edge Side Includes](#edge-side-includes)
   - [Regular expressions in conditions](#regular-expressions-in-conditions)
@@ -43,10 +44,10 @@ Using within your worker is as easy as including it and passing your request to
 import esi from "cloudflare-esi"
 
 export default {
-	async fetch(request: Request, env: any, ctx: any) {
-		const parser = new esi()
-		return parser.parse(request)
-	}
+   async fetch(request: Request, env: any, ctx: any) {
+      const parser = new esi()
+      return parser.parse(request)
+   }
 }
 
 ```
@@ -58,6 +59,7 @@ new esi(
    options?: // Config
    customESIFunction?: // Custom ESI Vars Function
    fetcher?: // Custom Fetch Function
+   postBodyFunction?: // Optional function that will be called once the body has been completed
 )
 ```
 
@@ -245,6 +247,10 @@ new esi(undefined, undefined, customFetcher)
 ```
 
 
+### Optional Post Body Function
+
+If you need to do extra work once the body has completed streaming, eg you record how many fetches your custom fetcher handles. Or need to fire off some context.waitUntil functions at the end you can fire them after this function has been called.
+
 
 ## Caching and upstream requests
 

src/index.ts

@@ -103,6 +103,7 @@ export type customESIVarsFunction = (
   request: Request,
 ) => Promise<customESIVars>;
 export type fetchFunction = (request: string | Request) => Promise<Response>;
+export type postBodyFunction = () => void | Promise<void>;
 
 const processorToken = "ESI";
 const processorVersion = 1.0;
@@ -111,19 +112,22 @@ export class esi {
   options: ESIConfig;
   esiFunction?: customESIVarsFunction;
   fetcher: fetchFunction;
+  postBodyFunction?: postBodyFunction;
 
   constructor(
     options?: ESIConfig,
     customESIFunction?: customESIVarsFunction,
     fetcher = fetch as fetchFunction,
+    postBodyFunction?: postBodyFunction,
   ) {
     const defaultConfig = {
       recursionLimit: 10,
       contentTypes: ["text/html", "text/plain"],
     };
     this.options = { ...defaultConfig, ...options };
     this.fetcher = fetcher;
-    if (customESIFunction) this.esiFunction = customESIFunction;
+    this.esiFunction = customESIFunction;
+    this.postBodyFunction = postBodyFunction;
   }
 
   async parse(origRequest: Request, recursion = 0): Promise<Response> {
@@ -252,7 +256,7 @@ export class esi {
     /**
      * Flushes output to the Writeable Stream
      */
-    async function flush_output() {
+    const flush_output = async () => {
       // Can't call this if we're waiting for an sync option to complete
       if (pending) {
         return;
@@ -287,8 +291,13 @@ export class esi {
       if (ended && output.length === 0) {
         const writer = writable.getWriter();
         await writer.close();
+
+        // we're completely done now notify the post body function
+        if (this.postBodyFunction) {
+          this.postBodyFunction();
+        }
       }
-    }
+    };
 
     const writer = (text: string, esi: boolean) => {
       if (esi) {

test/postBody.spec.ts

@@ -0,0 +1,83 @@
+import { checkSurrogate, testResponse, urlHandler } from "./helpers";
+import { esi, ESIConfig } from "../src";
+import { AddressInfo } from "net";
+import http from "http";
+
+const esiHead = {
+  "Content-Type": "text/html",
+  "Surrogate-Control": `content="ESI/1.0"`,
+};
+
+let parser: esi;
+let config: ESIConfig;
+const makeRequest = async function (request: string, details?: RequestInit) {
+  const reqUrl = new URL(request, `http://localhost:${port}`).toString();
+  const req = new Request(reqUrl, details);
+  return parser.parse(req);
+};
+
+const routeHandler = new urlHandler();
+// @ts-ignore
+const server = http.createServer(routeHandler.route);
+let port = 0;
+const testingDetails = {
+  port: 0,
+  hostname: "localhost",
+  proto: "http:",
+  url: `http://localhost:0`,
+};
+
+beforeAll(() => {
+  // Setup a basic HTTP server to handle traffic
+  server.listen(0);
+  const address = server.address() as AddressInfo;
+  port = address.port;
+  testingDetails.port = address.port;
+  testingDetails.url = `http://localhost:${port}`;
+});
+
+afterAll((done) => {
+  server.close(() => {
+    done();
+  });
+});
+
+test("TEST 1: postBody Function", async () => {
+  const url = `/post-body/test-1`;
+  let count = 0;
+  const postBody = function () {
+    expect(count).toEqual(3);
+    return;
+  };
+  parser = new esi(config, undefined, undefined, postBody);
+
+  const printFragment = function (
+    req: http.IncomingMessage,
+    res: testResponse,
+  ) {
+    count++;
+    const url = new URL(req.url as string, `http://localhost:${port}`);
+    const query = url.searchParams.get("a") ? url.searchParams.get("a") : "";
+    res.end(`FRAGMENT: ${query}\n`);
+  };
+  routeHandler.add(url, function (req, res) {
+    res.writeHead(200, esiHead);
+    res.say("1");
+    res.write(`<esi:include src="${url}/fragment_1" />`);
+    res.say("2");
+    res.write(`<esi:include src="${url}/fragment_1?a=2" />`);
+    res.write("3");
+    res.end(
+      `<esi:include src="http://localhost:${port}${url}/fragment_1?a=3" />`,
+    );
+  });
+  routeHandler.add(`${url}/fragment_1`, printFragment);
+  routeHandler.add(`${url}/fragment_1?a=2`, printFragment);
+  routeHandler.add(`${url}/fragment_1?a=3`, printFragment);
+  const res = await makeRequest(url);
+  expect(res.ok).toBeTruthy();
+  expect(checkSurrogate(res)).toBeTruthy();
+  expect(await res.text()).toEqual(
+    `1\nFRAGMENT: \n2\nFRAGMENT: 2\n3FRAGMENT: 3\n`,
+  );
+});