From c5bffb97cc03e1d78f0a180fbde43cbdaeb67dc3 Mon Sep 17 00:00:00 2001
From: Mike Cao <mike@mikecao.com>
Date: Sat, 22 Apr 2023 14:17:57 -0700
Subject: [PATCH] Improved error handling for useSession middleware.

---
 lib/middleware.ts | 15 ++++++++++-----
 lib/session.ts    |  6 +++---
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/lib/middleware.ts b/lib/middleware.ts
index b0785e2d..79c48404 100644
--- a/lib/middleware.ts
+++ b/lib/middleware.ts
@@ -20,14 +20,19 @@ export const useCors = createMiddleware(
 );
 
 export const useSession = createMiddleware(async (req, res, next) => {
-  const session = await findSession(req as NextApiRequestCollect);
+  try {
+    const session = await findSession(req as NextApiRequestCollect);
 
-  if (!session) {
-    log('useSession: Session not found');
-    return badRequest(res, 'Session not found.');
+    if (!session) {
+      log('useSession: Session not found');
+      return badRequest(res, 'Session not found.');
+    }
+
+    (req as any).session = session;
+  } catch (e: any) {
+    return badRequest(res, e.message);
   }
 
-  (req as any).session = session;
   next();
 });
 
diff --git a/lib/session.ts b/lib/session.ts
index af9f4c11..937bfef2 100644
--- a/lib/session.ts
+++ b/lib/session.ts
@@ -11,7 +11,7 @@ export async function findSession(req: NextApiRequestCollect) {
   const { payload } = getJsonBody<CollectRequestBody>(req);
 
   if (!payload) {
-    return null;
+    throw new Error('Invalid payload.');
   }
 
   // Check if cache token is passed
@@ -29,14 +29,14 @@ export async function findSession(req: NextApiRequestCollect) {
   const { website: websiteId, hostname, screen, language } = payload;
 
   if (!validate(websiteId)) {
-    return null;
+    throw new Error('Invalid website ID.');
   }
 
   // Find website
   const website = await loadWebsite(websiteId);
 
   if (!website) {
-    throw new Error(`Website not found: ${websiteId}`);
+    throw new Error(`Website not found: ${websiteId}.`);
   }
 
   const { userAgent, browser, os, ip, country, subdivision1, subdivision2, city, device } =