From 93f5bd30b9764e6911ed3034d14aeb4ef77067d8 Mon Sep 17 00:00:00 2001
From: HorizonCode <horizoncode88@gmail.com>
Date: Thu, 11 May 2023 09:32:18 +0200
Subject: [PATCH] add custom not found handler

---
 example/test.ts | 13 ++++++++
 http_server.ts  | 89 ++++++++++++++++++++++++++++++++++++-------------
 2 files changed, 79 insertions(+), 23 deletions(-)

diff --git a/example/test.ts b/example/test.ts
index f3c0533..2e3e6c6 100644
--- a/example/test.ts
+++ b/example/test.ts
@@ -3,6 +3,19 @@ import { HTTPServer } from "../http_server.ts";
 
 const httpServer = new HTTPServer();
 
+httpServer.error((req, _rep) => {
+  return JSON.stringify(
+    {
+      code: Status.NotFound,
+      message: "Route not found!",
+      path: req.path,
+      url: req.url
+    },
+    null,
+    2,
+  );
+})
+
 httpServer.add("GET", "/", (req, rep) => {
   rep.status(Status.Teapot)
     .header("working", "true")
diff --git a/http_server.ts b/http_server.ts
index 44f7d30..d60702f 100644
--- a/http_server.ts
+++ b/http_server.ts
@@ -28,6 +28,7 @@ export class HTTPServer {
   private routes = new Map<string, Route>();
   private staticLocalDir?: string;
   private staticServePath?: string;
+  private notFoundHandler?: RouteHandler;
 
   async listen(options: ListenOptions) {
     this.server = Deno.listen({
@@ -56,6 +57,8 @@ export class HTTPServer {
     for await (const requestEvent of httpConn) {
       const url = new URL(requestEvent.request.url);
       const filepath = decodeURIComponent(url.pathname);
+      const routeRequest = new RouteRequest(requestEvent.request);
+      const routeReply: RouteReply = new RouteReply();
 
       if (this.staticServePath && filepath.startsWith(this.staticServePath)) {
         const fileDir = filepath.split("/").slice(2).join("/");
@@ -69,17 +72,37 @@ export class HTTPServer {
           file = await Deno.open(pathLoc, { read: true });
         } catch {
           // If the file cannot be opened, return a "404 Not Found" response
-          await requestEvent.respondWith(
-            new Response(
-              JSON.stringify({
-                code: 404,
-                message: `File ${filepath} not found!`,
+          if (this.notFoundHandler) {
+            routeReply.status(Status.NotFound);
+            routeReply.type("application/json");
+            const notNoundHandle = await this.notFoundHandler(
+              routeRequest,
+              routeReply,
+            );
+            await requestEvent.respondWith(
+              new Response(notNoundHandle as string, {
+                status: routeReply.statusCode,
+                headers: routeReply.headers,
+                statusText: STATUS_TEXT[routeReply.statusCode],
               }),
-              {
-                status: Status.NotFound,
-              },
-            ),
-          );
+            );
+            continue;
+          } else {
+            await requestEvent.respondWith(
+              new Response(
+                JSON.stringify({
+                  code: 404,
+                  message: `File ${filepath} not found!`,
+                }),
+                {
+                  status: Status.NotFound,
+                  headers: {
+                    "Content-Type": "application/json",
+                  },
+                },
+              ),
+            );
+          }
           continue;
         }
 
@@ -88,8 +111,6 @@ export class HTTPServer {
         await requestEvent.respondWith(response);
         return;
       }
-      const routeRequest = new RouteRequest(requestEvent.request);
-      const routeReply: RouteReply = new RouteReply();
       const routeName = `${requestEvent.request.method}@${filepath}`;
       let route = this.routes.has(routeName)
         ? this.routes.get(routeName)
@@ -141,21 +162,43 @@ export class HTTPServer {
 
         continue;
       }
-
-      await requestEvent.respondWith(
-        new Response(
-          JSON.stringify({
-            code: 404,
-            message: `Route ${routeName} not found!`,
+      if (this.notFoundHandler) {
+        routeReply.status(Status.NotFound);
+        routeReply.type("application/json");
+        const notNoundHandle = await this.notFoundHandler(
+          routeRequest,
+          routeReply,
+        );
+        await requestEvent.respondWith(
+          new Response(notNoundHandle as string, {
+            status: routeReply.statusCode,
+            headers: routeReply.headers,
+            statusText: STATUS_TEXT[routeReply.statusCode],
           }),
-          {
-            status: Status.NotFound,
-          },
-        ),
-      );
+        );
+      } else {
+        await requestEvent.respondWith(
+          new Response(
+            JSON.stringify({
+              code: 404,
+              message: `Route ${routeName} not found!`,
+            }),
+            {
+              status: Status.NotFound,
+              headers: {
+                "Content-Type": "application/json",
+              },
+            },
+          ),
+        );
+      }
     }
   }
 
+  error(handler: RouteHandler) {
+    this.notFoundHandler = handler;
+  }
+
   get(path: string, handler: RouteHandler) {
     this.add("GET", path, handler);
   }