From 78e15684104f3ad3ddd6ddf4ac5839f1f598f653 Mon Sep 17 00:00:00 2001
From: Arkadiusz Sygulski <aareksio@gmail.com>
Date: Fri, 12 Jan 2024 18:41:42 +0700
Subject: [PATCH] fix(http, server): handle missing resources with http 404
 (#367)

Co-authored-by: Pooya Parsa <pooya@pi0.io>
---
 src/drivers/http.ts | 15 +++++++++++----
 src/server.ts       | 12 ++++++++++++
 test/server.test.ts |  9 +++++++++
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/src/drivers/http.ts b/src/drivers/http.ts
index 01fc9901..8ede6986 100644
--- a/src/drivers/http.ts
+++ b/src/drivers/http.ts
@@ -1,5 +1,5 @@
 import { defineDriver } from "./utils";
-import { $fetch as _fetch } from "ofetch";
+import { type FetchError, $fetch as _fetch } from "ofetch";
 import { joinURL } from "ufo";
 
 export interface HTTPOptions {
@@ -15,6 +15,13 @@ export default defineDriver((opts: HTTPOptions) => {
   const rBase = (key: string = "") =>
     joinURL(opts.base!, (key || "/").replace(/:/g, "/"), ":");
 
+  const catchFetchError = (error: FetchError, fallbackVal: any = null) => {
+    if (error?.response?.status === 404) {
+      return fallbackVal;
+    }
+    throw error;
+  };
+
   return {
     name: DRIVER_NAME,
     options: opts,
@@ -24,12 +31,12 @@ export default defineDriver((opts: HTTPOptions) => {
         headers: { ...opts.headers, ...topts.headers },
       })
         .then(() => true)
-        .catch(() => false);
+        .catch((err) => catchFetchError(err, false));
     },
     async getItem(key, tops = {}) {
       const value = await _fetch(r(key), {
         headers: { ...opts.headers, ...tops.headers },
-      });
+      }).catch(catchFetchError);
       return value;
     },
     async getItemRaw(key, topts) {
@@ -39,7 +46,7 @@ export default defineDriver((opts: HTTPOptions) => {
           ...opts.headers,
           ...topts.headers,
         },
-      });
+      }).catch(catchFetchError);
       return value;
     },
     async getMeta(key, topts) {
diff --git a/src/server.ts b/src/server.ts
index d7b0def4..60c832ec 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -76,11 +76,23 @@ export function createH3StorageHandler(
 
       const isRaw =
         getRequestHeader(event, "accept") === "application/octet-stream";
+
+      const checkNotFound = (value: any) => {
+        if (value === null) {
+          throw createError({
+            statusMessage: "KV value not found",
+            statusCode: 404,
+          });
+        }
+      };
+
       if (isRaw) {
         const value = await storage.getItemRaw(key);
+        checkNotFound(value);
         return value;
       } else {
         const value = await storage.getItem(key);
+        checkNotFound(value);
         return stringify(value);
       }
     }
diff --git a/test/server.test.ts b/test/server.test.ts
index adc30c1a..f5bee498 100644
--- a/test/server.test.ts
+++ b/test/server.test.ts
@@ -36,6 +36,15 @@ describe("server", () => {
     expect(await fetchStorage("foo/bar", { method: "DELETE" })).toBe("OK");
     expect(await fetchStorage("foo/bar/", {})).toMatchObject([]);
 
+    await expect(
+      fetchStorage("/non", { method: "GET" }).catch((error) => {
+        throw error.data;
+      })
+    ).rejects.toMatchObject({
+      statusCode: 404,
+      statusMessage: "KV value not found",
+    });
+
     await expect(
       fetchStorage("private/foo/bar", { method: "GET" }).catch((error) => {
         throw error.data;