Optimize header fixing logic and improve tests (#1036)

* Apply fixHeaders only once

* Don't fix cache-control headers if they've been changed

* update test

* linting

* changeset

* fix unit test

* Update packages/open-next/src/core/routing/util.ts

Co-authored-by: Victor Berchet <[email protected]>

* fixup! refactor

---------

Co-authored-by: Victor Berchet <[email protected]>

Author: conico974

.changeset/large-pants-press.md

@@ -0,0 +1,5 @@
+---
+"@opennextjs/aws": patch
+---
+
+Fix cache-control header set in middleware being overriden for ISR route

packages/open-next/src/core/routing/util.ts

@@ -356,25 +356,31 @@ export async function revalidateIfRequired(
  * @__PURE__
  */
 export function fixISRHeaders(headers: OutgoingHttpHeaders) {
+  const sMaxAgeRegex = /s-maxage=(\d+)/;
+  const match = headers[CommonHeaders.CACHE_CONTROL]?.match(sMaxAgeRegex);
+  const sMaxAge = match ? Number.parseInt(match[1]) : undefined;
+  // We only apply the fix if the cache-control header contains s-maxage
+  if (!sMaxAge) {
+    return;
+  }
   if (headers[CommonHeaders.NEXT_CACHE] === "REVALIDATED") {
     headers[CommonHeaders.CACHE_CONTROL] =
       "private, no-cache, no-store, max-age=0, must-revalidate";
     return;
   }
   const _lastModified = globalThis.__openNextAls.getStore()?.lastModified ?? 0;
   if (headers[CommonHeaders.NEXT_CACHE] === "HIT" && _lastModified > 0) {
-    // calculate age
-    const age = Math.round((Date.now() - _lastModified) / 1000);
-    // extract s-maxage from cache-control
-    const regex = /s-maxage=(\d+)/;
-    const cacheControl = headers[CommonHeaders.CACHE_CONTROL];
-    debug("cache-control", cacheControl, _lastModified, Date.now());
-    if (typeof cacheControl !== "string") return;
-    const match = cacheControl.match(regex);
-    const sMaxAge = match ? Number.parseInt(match[1]) : undefined;
+    debug(
+      "cache-control",
+      headers[CommonHeaders.CACHE_CONTROL],
+      _lastModified,
+      Date.now(),
+    );
 
     // 31536000 is the default s-maxage value for SSG pages
     if (sMaxAge && sMaxAge !== 31536000) {
+      // calculate age
+      const age = Math.round((Date.now() - _lastModified) / 1000);
       const remainingTtl = Math.max(sMaxAge - age, 1);
       headers[CommonHeaders.CACHE_CONTROL] =
         `s-maxage=${remainingTtl}, stale-while-revalidate=2592000`;

packages/open-next/src/http/openNextResponse.ts

@@ -22,6 +22,7 @@ export class OpenNextNodeResponse extends Transform implements ServerResponse {
   headers: OutgoingHttpHeaders = {};
   headersSent = false;
   _chunks: Buffer[] = [];
+  headersAlreadyFixed = false;
 
   private _cookies: string[] = [];
   private responseStream?: Writable;
@@ -69,7 +70,7 @@ export class OpenNextNodeResponse extends Transform implements ServerResponse {
   }
 
   constructor(
-    private fixHeaders: (headers: OutgoingHttpHeaders) => void,
+    private fixHeadersFn: (headers: OutgoingHttpHeaders) => void,
     private onEnd: (headers: OutgoingHttpHeaders) => Promise<void>,
     private streamCreator?: StreamCreator,
     private initialHeaders?: OutgoingHttpHeaders,
@@ -271,6 +272,14 @@ export class OpenNextNodeResponse extends Transform implements ServerResponse {
    * OpenNext specific method
    */
 
+  fixHeaders(headers: OutgoingHttpHeaders) {
+    if (this.headersAlreadyFixed) {
+      return;
+    }
+    this.fixHeadersFn(headers);
+    this.headersAlreadyFixed = true;
+  }
+
   getFixedHeaders(): OutgoingHttpHeaders {
     // Do we want to apply this on writeHead?
     this.fixHeaders(this.headers);

packages/tests-e2e/tests/appRouter/isr.test.ts

@@ -45,13 +45,26 @@ test("headers", async ({ page }) => {
     return response.status() === 200;
   });
   await page.goto("/isr");
+  let hasBeenStale = false;
+  let hasBeenHit = false;
 
   while (true) {
     const response = await responsePromise;
     const headers = response.headers();
+    expect(headers["cache-control"]).toBe(
+      "max-age=10, stale-while-revalidate=999",
+    );
+    const cacheHeader =
+      headers["x-nextjs-cache"] ?? headers["x-opennext-cache"];
+    if (cacheHeader === "STALE") {
+      hasBeenStale = true;
+    }
+    if (cacheHeader === "HIT") {
+      hasBeenHit = true;
+    }
 
     // this was set in middleware
-    if (headers["cache-control"] === "max-age=10, stale-while-revalidate=999") {
+    if (hasBeenStale && hasBeenHit) {
       break;
     }
     await page.waitForTimeout(1000);

packages/tests-unit/tests/core/routing/util.test.ts

@@ -748,6 +748,7 @@ describe("fixISRHeaders", () => {
   it("should set cache-control directive to must-revalidate when x-nextjs-cache is REVALIDATED", () => {
     const headers: Record<string, string> = {
       "x-nextjs-cache": "REVALIDATED",
+      "cache-control": "s-maxage=10, stale-while-revalidate=86400",
     };
     fixISRHeaders(headers);
 
@@ -771,13 +772,33 @@ describe("fixISRHeaders", () => {
   it("should set cache-control directive to stale-while-revalidate when x-nextjs-cache is STALE", () => {
     const headers: Record<string, string> = {
       "x-nextjs-cache": "STALE",
+      "cache-control": "s-maxage=86400", // 1 day
     };
     fixISRHeaders(headers);
 
     expect(headers["cache-control"]).toBe(
       "s-maxage=2, stale-while-revalidate=2592000",
     );
   });
+
+  it("should not modify cache-control when cache-control is missing", () => {
+    const headers: Record<string, string> = {
+      "x-nextjs-cache": "HIT",
+    };
+    fixISRHeaders(headers);
+
+    expect(headers["cache-control"]).toBeUndefined();
+  });
+
+  it("should not modify cache-control when cache-control has no s-maxage", () => {
+    const headers: Record<string, string> = {
+      "x-nextjs-cache": "HIT",
+      "cache-control": "private, max-age=0",
+    };
+    fixISRHeaders(headers);
+
+    expect(headers["cache-control"]).toBe("private, max-age=0");
+  });
 });
 
 describe("invalidateCDNOnRequest", () => {