From 06820b5ae02bae3403ed233e45bf66dd8d1726ed Mon Sep 17 00:00:00 2001
From: Pooya Parsa <pooya@pi0.io>
Date: Wed, 18 Oct 2023 00:35:33 +0200
Subject: [PATCH] fix(server): improve 304 handling

closes #178
---
 src/server.ts | 43 ++++++++++++++++++++++++-------------------
 1 file changed, 24 insertions(+), 19 deletions(-)

diff --git a/src/server.ts b/src/server.ts
index 696153d..95decd1 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -13,6 +13,7 @@ import {
   createError,
   H3Event,
   H3Error,
+  send,
 } from "h3";
 import { IPX } from "./ipx";
 
@@ -79,18 +80,10 @@ export function createIPXH3Handler(ipx: IPX) {
     // Get image meta from source
     const sourceMeta = await img.getSourceMeta();
 
-    // Caching headers
-    if (sourceMeta.mtime) {
-      if (
-        getRequestHeader(event, "if-modified-since") &&
-        new Date(getRequestHeader(event, "if-modified-since") || "") >=
-          sourceMeta.mtime
-      ) {
-        setResponseStatus(event, 304);
-        return null;
-      }
-      setResponseHeader(event, "last-modified", sourceMeta.mtime.toUTCString());
-    }
+    // Send CSP headers to prevent XSS
+    setResponseHeader(event, "content-security-policy", "default-src 'none'");
+
+    // Send Cache-Control header
     if (typeof sourceMeta.maxAge === "number") {
       setResponseHeader(
         event,
@@ -99,25 +92,37 @@ export function createIPXH3Handler(ipx: IPX) {
       );
     }
 
-    // Get converted image
+    // Handle modified time if available
+    if (sourceMeta.mtime) {
+      // Send Last-Modified header
+      setResponseHeader(event, "last-modified", sourceMeta.mtime.toUTCString());
+
+      // Check for last-modified request header
+      const _ifModifiedSince = getRequestHeader(event, "if-modified-since");
+      if (_ifModifiedSince && new Date(_ifModifiedSince) >= sourceMeta.mtime) {
+        setResponseStatus(event, 304);
+        return send(event);
+      }
+    }
+
+    // Process image
     const { data, format } = await img.process();
 
-    // ETag
+    // Generate and send ETag header
     const etag = getEtag(data);
     setResponseHeader(event, "etag", etag);
+
+    // Check for if-none-match request header
     if (etag && getRequestHeader(event, "if-none-match") === etag) {
       setResponseStatus(event, 304);
-      return null;
+      return send(event);
     }
 
-    // Mime
+    // Content-Type header
     if (format) {
       setResponseHeader(event, "content-type", `image/${format}`);
     }
 
-    // Prevent XSS
-    setResponseHeader(event, "content-security-policy", "default-src 'none'");
-
     return data;
   };