refactor(astro): streamline cookie handling and improve type definitions in index.ts

Author: RyanTamulevicz

src/astro/index.ts

@@ -1,37 +1,28 @@
 import { betterFetch } from "@better-fetch/fetch";
-import type { betterAuth } from "better-auth";
+import { betterAuth } from "better-auth";
 import { createCookieGetter } from "better-auth/cookies";
+import type { APIContext, AstroCookies } from "astro";
 import { ConvexHttpClient } from "convex/browser";
 import type {
   FunctionReference,
   FunctionReturnType,
   GenericActionCtx,
   GenericDataModel,
 } from "convex/server";
-import type { APIContext, AstroCookies } from "astro";
-import { type CreateAuth, getStaticAuth } from "../client";
 import { JWT_COOKIE_NAME } from "../plugins/convex";
+import { type CreateAuth, getStaticAuth } from "../client";
 
-type CookieReader = (name: string) => string | undefined;
-type CookieSource = AstroCookies | CookieReader | undefined;
 type AstroRequestContext =
   | Pick<APIContext, "request" | "cookies">
-  | {
-      request: Request;
-      cookies?: CookieSource;
-    };
-type CookieInput =
-  | CookieSource
+  | { request: Request; cookies?: AstroCookies };
+
+type CookieSource =
   | AstroRequestContext
+  | AstroCookies
   | Request
   | Headers
   | string
-  | {
-      request?: Request;
-      headers?: Headers | string;
-      cookies?: CookieSource | string;
-      cookie?: string;
-    };
+  | undefined;
 
 const safeDecode = (value: string) => {
   try {
@@ -41,127 +32,74 @@ const safeDecode = (value: string) => {
   }
 };
 
-const cookieReaderFromString = (cookieHeader: string): CookieReader => {
+const isAstroCookies = (value: unknown): value is AstroCookies =>
+  !!value &&
+  typeof value === "object" &&
+  "get" in value &&
+  typeof (value as { get?: unknown }).get === "function";
+
+const readCookieFromHeader = (
+  cookieHeader: string | null | undefined,
+  name: string
+) => {
   if (!cookieHeader) {
-    return () => undefined;
+    return undefined;
   }
-  const pairs = cookieHeader.split(/;\s*/).filter(Boolean);
-  const store = new Map<string, string>();
-  for (const pair of pairs) {
+  for (const pair of cookieHeader.split(/;\s*/)) {
+    if (!pair) {
+      continue;
+    }
     const separatorIndex = pair.indexOf("=");
     if (separatorIndex === -1) {
       continue;
     }
     const key = safeDecode(pair.slice(0, separatorIndex).trim());
-    const value = safeDecode(pair.slice(separatorIndex + 1));
-    if (!store.has(key)) {
-      store.set(key, value);
-    }
-  }
-  if (store.size === 0) {
-    return () => undefined;
-  }
-  return (name) => {
-    if (!name) {
-      return undefined;
+    if (key !== name) {
+      continue;
     }
-    return store.get(name);
-  };
-};
-
-const isAstroCookies = (source: unknown): source is AstroCookies => {
-  if (!source || typeof source !== "object") {
-    return false;
-  }
-  return (
-    "get" in source &&
-    typeof (source as { get?: unknown }).get === "function" &&
-    "has" in source &&
-    typeof (source as { has?: unknown }).has === "function" &&
-    "merge" in source &&
-    typeof (source as { merge?: unknown }).merge === "function"
-  );
-};
-
-const isHeadersLike = (source: unknown): source is Headers => {
-  return typeof Headers !== "undefined" && source instanceof Headers;
-};
-
-const isRequestLike = (source: unknown): source is Request => {
-  if (typeof Request !== "undefined" && source instanceof Request) {
-    return true;
+    const rawValue = pair.slice(separatorIndex + 1);
+    return safeDecode(rawValue);
   }
-  if (!source || typeof source !== "object") {
-    return false;
-  }
-  return (
-    "headers" in source &&
-    typeof (source as { headers?: unknown }).headers !== "undefined" &&
-    "method" in source
-  );
+  return undefined;
 };
 
-const normalizeCookieSource = (source: CookieInput): CookieSource => {
+const readCookie = (source: CookieSource, name: string): string | undefined => {
   if (!source) {
     return undefined;
   }
-  if (typeof source === "function") {
-    return source;
-  }
-  if (isAstroCookies(source)) {
-    return source;
+  if (typeof source === "string") {
+    return readCookieFromHeader(source, name);
   }
-  if (isRequestLike(source)) {
-    return normalizeCookieSource((source as Request).headers);
+  if (source instanceof Headers) {
+    return readCookieFromHeader(source.get("cookie"), name);
   }
-  if (isHeadersLike(source)) {
-    const cookieHeader = (source as Headers).get("cookie") ?? "";
-    return cookieHeader ? cookieReaderFromString(cookieHeader) : undefined;
+  if (source instanceof Request) {
+    return readCookie(source.headers, name);
   }
-  if (typeof source === "string") {
-    return cookieReaderFromString(source);
+  if (isAstroCookies(source)) {
+    const cookie = source.get(name);
+    return typeof cookie?.value === "string" ? cookie.value : undefined;
   }
-  if (typeof source === "object") {
-    if ("cookies" in source && source.cookies) {
-      const normalized = normalizeCookieSource(source.cookies as CookieInput);
-      if (normalized) {
-        return normalized;
-      }
-    }
-    if ("request" in source && source.request) {
-      const normalized = normalizeCookieSource(source.request as CookieInput);
-      if (normalized) {
-        return normalized;
-      }
-    }
-    if ("headers" in source && source.headers) {
-      const normalized = normalizeCookieSource(source.headers as CookieInput);
-      if (normalized) {
-        return normalized;
-      }
+  if (
+    typeof source === "object" &&
+    source !== null &&
+    "cookies" in source &&
+    source.cookies
+  ) {
+    const fromStore = readCookie(source.cookies, name);
+    if (fromStore) {
+      return fromStore;
     }
-    if ("cookie" in source && typeof source.cookie === "string") {
-      const normalized = normalizeCookieSource(source.cookie);
-      if (normalized) {
-        return normalized;
-      }
-    }
-  }
-  return undefined;
-};
-
-const createCookieReader = (source: CookieInput): CookieReader => {
-  const normalized = normalizeCookieSource(source);
-  if (typeof normalized === "function") {
-    return normalized;
   }
-  if (normalized && typeof normalized.get === "function") {
-    return (name) => {
-      const cookie = normalized.get(name) as { value?: string } | undefined;
-      return typeof cookie?.value === "string" ? cookie.value : undefined;
-    };
+  if (
+    typeof source === "object" &&
+    source !== null &&
+    "request" in source &&
+    source.request
+  ) {
+    return readCookie(source.request, name);
   }
-  return () => undefined;
+  return undefined;
 };
 
 export const getCookieName = <DataModel extends GenericDataModel>(
@@ -174,24 +112,23 @@ export const getCookieName = <DataModel extends GenericDataModel>(
 
 export const getToken = <DataModel extends GenericDataModel>(
   createAuth: CreateAuth<DataModel>,
-  cookies?: CookieInput
+  cookies?: CookieSource
 ) => {
   const sessionCookieName = getCookieName(createAuth);
-  const readCookie = createCookieReader(cookies);
-  const token = readCookie(sessionCookieName);
+  const token = readCookie(cookies, sessionCookieName);
 
   if (!token) {
     const isSecure = sessionCookieName.startsWith("__Secure-");
     const insecureCookieName = sessionCookieName.replace("__Secure-", "");
     const secureCookieName = isSecure
       ? sessionCookieName
       : `__Secure-${insecureCookieName}`;
-    const secureToken = readCookie(secureCookieName);
-    const insecureToken = readCookie(insecureCookieName);
+    const secureToken = readCookie(cookies, secureCookieName);
+    const insecureToken = readCookie(cookies, insecureCookieName);
 
     if (isSecure && insecureToken) {
       console.warn(
-        `Looking for secure cookie ${sessionCookieName} but found insecure cookie ${sessionCookieName.replace("__Secure-", "")}`
+        `Looking for secure cookie ${sessionCookieName} but found insecure cookie ${insecureCookieName}`
       );
     }
     if (!isSecure && secureToken) {
@@ -206,18 +143,16 @@ export const getToken = <DataModel extends GenericDataModel>(
 
 export const setupFetchClient = async <DataModel extends GenericDataModel>(
   createAuth: CreateAuth<DataModel>,
-  cookies?: CookieInput,
+  cookies?: CookieSource,
   opts?: { convexUrl?: string }
 ) => {
-  const readCookie = createCookieReader(cookies);
   const createClient = () => {
     const convexUrl = opts?.convexUrl ?? process.env.VITE_CONVEX_URL;
     if (!convexUrl) {
       throw new Error("VITE_CONVEX_URL is not set");
     }
-    const sessionCookieName = getCookieName(createAuth);
-    const token = readCookie(sessionCookieName);
     const client = new ConvexHttpClient(convexUrl);
+    const token = getToken(createAuth, cookies);
     if (token) {
       client.setAuth(token);
     }
@@ -295,9 +230,7 @@ export const getAuth = async <DataModel extends GenericDataModel>(
   if (!request) {
     throw new Error("No request found");
   }
-  const readCookie = createCookieReader(context);
-  const sessionCookieName = getCookieName(createAuth);
-  const token = readCookie(sessionCookieName);
+  const token = getToken(createAuth, context);
   const { session } = await fetchSession(request, opts);
   return {
     userId: session?.user.id,