Split logic for findSession.

lib/cache.ts

@@ -1,4 +1,4 @@
-import { User, Website, Team } from '@prisma/client';
+import { User, Website, Team, TeamUser } from '@prisma/client';
 import redis from '@umami/redis-client';
 import { lightFormat, startOfMonth } from 'date-fns';
 import { getAllWebsitesByUser, getSession, getUser, getViewTotals, getWebsite } from '../queries';
@@ -35,10 +35,10 @@ async function deleteObject(key, soft = false) {
 
 async function fetchWebsite(id): Promise<
   Website & {
-    team?: Team;
+    team?: Team & { teamUsers: TeamUser[] };
   }
 > {
-  return fetchObject(`website:${id}`, () => getWebsite({ id }));
+  return fetchObject(`website:${id}`, () => getWebsite({ id }, true));
 }
 
 async function storeWebsite(data) {

lib/middleware.ts

@@ -1,20 +1,57 @@
-import { createMiddleware, unauthorized, badRequest, parseSecureToken } from 'next-basics';
-import debug from 'debug';
-import cors from 'cors';
-import { validate } from 'uuid';
 import redis from '@umami/redis-client';
-import { findSession } from 'lib/session';
+import cors from 'cors';
+import debug from 'debug';
 import { getAuthToken, parseShareToken } from 'lib/auth';
-import { secret } from 'lib/crypto';
 import { ROLES } from 'lib/constants';
+import { secret } from 'lib/crypto';
+import { isOverApiLimit, findSession, findWebsite, useSessionCache } from 'lib/session';
+import {
+  badRequest,
+  createMiddleware,
+  parseSecureToken,
+  tooManyRequest,
+  unauthorized,
+} from 'next-basics';
+import { NextApiRequestCollect } from 'pages/api/collect';
+import { validate } from 'uuid';
 import { getUser } from '../queries';
+import { getJsonBody } from './detect';
 
 const log = debug('umami:middleware');
 
 export const useCors = createMiddleware(cors());
 
-export const useSession = createMiddleware(async (req: any, res, next) => {
-  const session = await findSession(req);
+export const useSession = createMiddleware(async (req: NextApiRequestCollect, res, next) => {
+  // Verify payload
+  const { payload } = getJsonBody(req);
+
+  const { website: websiteId } = payload;
+
+  if (!payload) {
+    log('useSession: No payload');
+    return badRequest(res);
+  }
+
+  // Get session from cache
+  let session = await useSessionCache(req);
+
+  // Get or create session
+  if (!session) {
+    const website = await findWebsite(websiteId);
+
+    if (!website) {
+      log('useSession: Website not found');
+      return badRequest(res);
+    }
+
+    if (process.env.ENABLE_COLLECT_LIMIT) {
+      if (isOverApiLimit(website)) {
+        return tooManyRequest(res, 'Collect currently exceeds monthly limit of 10000.');
+      }
+    }
+
+    session = await findSession(req, payload);
+  }
 
   if (!session) {
     log('useSession: Session not found');

lib/session.ts

@@ -2,17 +2,14 @@ import { Session, Team, Website } from '@prisma/client';
 import cache from 'lib/cache';
 import clickhouse from 'lib/clickhouse';
 import { secret, uuid } from 'lib/crypto';
-import { getClientInfo, getJsonBody } from 'lib/detect';
+import { getClientInfo } from 'lib/detect';
 import { parseToken } from 'next-basics';
 import { createSession, getSession, getWebsite } from 'queries';
-import { validate } from 'uuid';
 
-export async function findSession(req): Promise<{
-  error?: {
-    status: number;
-    message: string;
-  };
-  session?: {
+export async function findSession(
+  req,
+  payload,
+): Promise<{
   id: string;
   websiteId: string;
   hostname: string;
@@ -22,53 +19,15 @@ export async function findSession(req): Promise<{
   screen: string;
   language: string;
   country: string;
-  };
-  website?: Website & { team?: Team };
 }> {
-  const { payload } = getJsonBody(req);
-
-  if (!payload) {
-    return null;
-  }
-
-  // Verify payload
   const { website: websiteId, hostname, screen, language } = payload;
 
-  // Find website
-  let website: Website & { team?: Team } = null;
-
-  if (cache.enabled) {
-    website = await cache.fetchWebsite(websiteId);
-  } else {
-    website = await getWebsite({ id: websiteId });
-  }
-
-  if (!website || website.deletedAt) {
-    throw new Error(`Website not found: ${websiteId}`);
-  }
-
-  // Check if cache token is passed
-  const cacheToken = req.headers['x-umami-cache'];
-
-  if (cacheToken) {
-    const result = await parseToken(cacheToken, secret());
-
-    if (result) {
-      return { session: result, website };
-    }
-  }
-
-  if (!validate(websiteId)) {
-    return null;
-  }
-
   const { userAgent, browser, os, ip, country, device } = await getClientInfo(req, payload);
   const sessionId = uuid(websiteId, hostname, ip, userAgent);
 
   // Clickhouse does not require session lookup
   if (clickhouse.enabled) {
     return {
-      session: {
       id: sessionId,
       websiteId,
       hostname,
@@ -78,8 +37,6 @@ export async function findSession(req): Promise<{
       screen,
       language,
       country,
-      },
-      website,
     };
   }
 
@@ -113,5 +70,61 @@ export async function findSession(req): Promise<{
     }
   }
 
-  return { session, website };
+  return session;
+}
+
+export async function useSessionCache(req: any): Promise<{
+  id: string;
+  websiteId: string;
+  hostname: string;
+  browser: string;
+  os: string;
+  device: string;
+  screen: string;
+  language: string;
+  country: string;
+}> {
+  // Check if cache token is passed
+  const cacheToken = req.headers['x-umami-cache'];
+
+  if (cacheToken) {
+    const result = await parseToken(cacheToken, secret());
+
+    if (result) {
+      return result;
+    }
+  }
+
+  return null;
+}
+
+export async function findWebsite(websiteId: string) {
+  let website: Website & { team?: Team } = null;
+
+  if (cache.enabled) {
+    website = await cache.fetchWebsite(websiteId);
+  } else {
+    website = await getWebsite({ id: websiteId }, true);
+  }
+
+  if (!website || website.deletedAt) {
+    throw new Error(`Website not found: ${websiteId}`);
+  }
+
+  return website;
+}
+
+export async function isOverApiLimit(website) {
+  const userId = website.userId ? website.userId : website.team.teamu.userId;
+
+  const limit = await cache.fetchCollectLimit(userId);
+
+  // To-do: Need to implement logic to find user-specific limit. Defaulted to 10k.
+  if (limit > 10000) {
+    return true;
+  }
+
+  await cache.incrementCollectLimit(userId);
+
+  return false;
 }

pages/api/collect.ts

@@ -1,28 +1,14 @@
 const { Resolver } = require('dns').promises;
-import isbot from 'isbot';
 import ipaddr from 'ipaddr.js';
-import {
-  createToken,
-  unauthorized,
-  send,
-  badRequest,
-  forbidden,
-  tooManyRequest,
-} from 'next-basics';
-import { savePageView, saveEvent } from 'queries';
-import { useCors, useSession } from 'lib/middleware';
-import { getJsonBody, getIpAddress } from 'lib/detect';
+import isbot from 'isbot';
 import { secret } from 'lib/crypto';
+import { getIpAddress, getJsonBody } from 'lib/detect';
+import { useCors, useSession } from 'lib/middleware';
 import { NextApiRequest, NextApiResponse } from 'next';
-import cache from 'lib/cache';
-import { Team, Website } from '@prisma/client';
+import { badRequest, createToken, forbidden, send, unauthorized } from 'next-basics';
+import { saveEvent, savePageView } from 'queries';
 
 export interface NextApiRequestCollect extends NextApiRequest {
-  session: {
-    error?: {
-      status: number;
-      message: string;
-    };
   session?: {
     id: string;
     websiteId: string;
@@ -34,8 +20,6 @@ export interface NextApiRequestCollect extends NextApiRequest {
     language: string;
     country: string;
   };
-    website?: Website & { team?: Team };
-  };
 }
 
 export default async (req: NextApiRequestCollect, res: NextApiResponse) => {
@@ -104,21 +88,7 @@ export default async (req: NextApiRequestCollect, res: NextApiResponse) => {
 
   await useSession(req, res);
 
-  const { session, website } = req.session;
-
-  // Check collection limit
-  if (process.env.ENABLE_COLLECT_LIMIT) {
-    const userId = website.userId ? website.userId : website.team.userId;
-
-    const limit = await cache.fetchCollectLimit(userId);
-
-    // To-do: Need to implement logic to find user-specific limit. Defaulted to 10k.
-    if (limit > 10000) {
-      return tooManyRequest(res, 'Collect currently exceeds monthly limit of 10000.');
-    }
-
-    await cache.incrementCollectLimit(userId);
-  }
+  const session = req.session;
 
   if (process.env.REMOVE_TRAILING_SLASH) {
     url = url.replace(/\/$/, '');

queries/admin/website.ts

@@ -1,18 +1,29 @@
-import { Prisma, Team, Website } from '@prisma/client';
+import { Prisma, Team, TeamUser, Website } from '@prisma/client';
 import cache from 'lib/cache';
 import prisma from 'lib/prisma';
 import { runQuery, CLICKHOUSE, PRISMA } from 'lib/db';
 
-export async function getWebsite(where: Prisma.WebsiteWhereUniqueInput): Promise<
+export async function getWebsite(
+  where: Prisma.WebsiteWhereUniqueInput,
+  includeTeamData = false,
+): Promise<
   Website & {
-    team?: Team;
+    team?: Team & { teamUsers: TeamUser[] };
   }
 > {
+  prisma.client.team.findMany();
+
   return prisma.client.website.findUnique({
     where,
+    include: includeTeamData
+      ? {
+          team: {
             include: {
-      team: true,
+              teamUsers: true,
             },
+          },
+        }
+      : {},
   });
 }