1 files changed, 17 insertions, 14 deletions
diff --git a/src/index.ts b/src/index.ts
index 8ae1d0d..e524527 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -19,6 +19,7 @@ function getFilename(name: string): string {
 
 type File = {
   filename: string;
+  headers: Record<string, string>;
   size: number;
   mtime: string;
 };
@@ -33,6 +34,10 @@ function registerFile(pathname: string, absPath: string, stat: Stats): void {
   }
   files.set(pathname, {
     filename: absPath,
+    headers:
+      pathname === "/404.html"
+        ? Object.assign({}, defaultHeaders, { "cache-control": "no-cache" })
+        : defaultHeaders,
     size: stat.size,
     mtime: stat.mtime.toUTCString(),
   });
@@ -61,23 +66,21 @@ function walkDirectory(root: string, dir: string) {
 
 walkDirectory(publicDir, "");
 
+async function serveFile(file: File | undefined): Promise<Response> {
+  if (file && (await fsp.exists(file.filename))) {
+    return new Response(Bun.file(file.filename), {
+      headers: file.headers,
+      status: 200,
+    });
+  } else {
+    return serveFile(files.get("/404.html"));
+  }
+}
+
 export default withHtmlLiveReload({
   fetch: async function (request) {
     const pathname = new URL(request.url).pathname;
     const file = files.get(pathname);
-    if (file) {
-      return new Response(Bun.file(file.filename), {
-        headers: defaultHeaders,
-        status: 200,
-      });
-    }
-
-    return new Response(Bun.file(getFilename("404.html")), {
-      headers: Object.assign({}, defaultHeaders, {
-        "cache-control": "max-age=5, no-cache",
-      }),
-      status: 404,
-      statusText: "Not Found",
-    });
+    return serveFile(file);
   },
 });