feat(billing): credit handling

docker-compose.yml

@@ -69,6 +69,18 @@ services:
     networks:
       - vulnscanner
 
+  global-service:
+    build:
+      dockerfile: ./global.dockerfile
+      target: dev
+      context: ./
+
+    image: ninjeneer/vuln-scanner-global-service
+    volumes:
+      - .:/app
+    networks:
+      - vulnscanner
+
 networks:
   vulnscanner:
     name: vulnscanner

global.dockerfile

@@ -0,0 +1,21 @@
+FROM node:alpine as builder
+WORKDIR /app
+COPY --chown=node:node package.json yarn.lock ./
+RUN yarn install --frozen-lockfile 
+COPY --chown=node:node . .
+RUN yarn build
+
+
+# Dev stage
+FROM builder as dev
+WORKDIR /app
+EXPOSE 3000
+CMD ["yarn", "dev:global"]
+
+
+# Prod stage
+FROM node:alpine as prod
+WORKDIR /app
+COPY --from=builder /app/node_modules ./dist/node_modules
+COPY --from=builder /app/dist/src ./dist/src
+CMD ["node", "dist/src/services/global/server.js"]

package.json

@@ -9,6 +9,7 @@
     "dev:jobs": "nodemon src/services/jobs/index.ts",
     "dev:billing": "nodemon src/services/billing/server.ts",
     "dev:stats": "nodemon src/services/stats/server.ts",
+    "dev:global": "nodemon src/services/global/index.ts",
     "build": "swc src -d dist/src",
     "test": "jest test"
   },

src/config.ts

@@ -11,4 +11,10 @@ export const getConfig = () => {
     } : {
         ...commonConfig
     }
+}
+
+export const creditsPlanMapping = {
+    'free': 10,
+    'premium': 100,
+    'enterprise': 1000
 }

src/exceptions/exceptions.ts

@@ -28,6 +28,12 @@ export class UserDoesNotExist extends Error {
     }
 }
 
+export class UserHasNotEnoughCredits extends Error {
+    constructor(id = "") {
+        super(`User id ${id} has not enough credits`)
+    }
+}
+
 export class StripePriceDoesNotExist extends Error {
     constructor(id = "") {
         super(`Stripe Price ${id} does not exist`)

src/models/settings.ts

@@ -1,3 +1,5 @@
 export type UserSettings = Partial<{
+    id: string
+    userId: string
     plan: string
 }>

src/services/billing/DAL/stripe/stripeTypes.ts

@@ -128,5 +128,6 @@ type Automatictax = {
 export type StripeProduct = Stripe.Product & {
     metadata: {
         plan?: string
+        credits?: number
     }
 }

src/services/billing/billing.ts

@@ -1,5 +1,6 @@
-import { UserDoesNotExist } from '../../exceptions/exceptions'
+import { MissingData, UserDoesNotExist } from '../../exceptions/exceptions'
 import { User } from '../../models/user'
+import { updateUserCredits } from '../../storage/credits.storage'
 import { getUserById, getUserSettings, updateUserSettings } from '../../storage/user.storage'
 import * as billing from './DAL/stripe/stripe'
 import { StripeCheckoutSessionCompleted } from './DAL/stripe/stripeTypes'
@@ -26,14 +27,15 @@ export const handleEvent = (eventType: string, data: any) => {
 }
 
 const handleCheckoutSession = async (data: StripeCheckoutSessionCompleted['data']) => {
-    const { userId, plan } = data.object.metadata
+    const { userId, plan, credits } = data.object.metadata
     if (!userId) {
-        console.error('Missing userId in billing hook event')
-        return
+        throw new MissingData('userId', 'billing plan hook event')
     }
     if (!plan) {
-        console.error('Missing plan in billing hook event')
-        return
+        throw new MissingData('plan', 'billing hook event')
+    }
+    if (!credits) {
+        throw new MissingData('credits', 'billing hook event')
     }
 
     const user = await getUserById(userId)
@@ -49,4 +51,7 @@ const handleCheckoutSession = async (data: StripeCheckoutSessionCompleted['data'
         await updateUserSettings(userId, { ...userSettings, plan })
         console.log(`[BILLING][CHECKOUT][NEW SUB] Updated settings of user ${user.email}`)
     }
+
+    // Give the user the amount of credit subscribed for
+    await updateUserCredits(userId, credits)
 }

src/services/global/events.ts

@@ -0,0 +1,32 @@
+import supabase from "../../storage/supabase"
+import {listenSettings} from '../../storage/settings.storage'
+import { MissingData } from "../../exceptions/exceptions"
+import { updateUserCredits } from "../../storage/credits.storage"
+import { creditsPlanMapping } from "../../config"
+
+export const listenEvents = () => {
+    listenSettingsForFreePlanCredits()
+}
+
+const listenSettingsForFreePlanCredits = () => {
+    console.log('[EVENT][SETTINGS] Listening for user settings events...')
+    listenSettings(async (change) => {
+        // New user has been created
+        if (change.eventType === 'INSERT') {
+            const { id, userId, plan } = change.new
+            if (!userId) {
+                throw new MissingData('userId', `user settings ${id}`)
+            }
+            if (!plan) {
+                throw new MissingData('plan', `user settings ${id}`)
+            }
+
+            // Credits amount are handled by Stripe webhook, but we need to handle Free plan registration
+            if (plan === 'free') {
+                const credits = creditsPlanMapping[plan]
+                await updateUserCredits(userId, credits)
+                console.log(`[EVENTS][SETTINGS] Set ${credits} credits for FREE plan to ${userId}`)
+            }
+        }
+    })
+}

src/services/global/index.ts

@@ -0,0 +1,3 @@
+import { listenEvents } from "./events";
+
+listenEvents()

src/services/requests/scanService.ts

@@ -4,17 +4,24 @@ import * as scanStorage from "../../storage/scan.storage";
 import { ScanStatus, ScanRequestResponse, ScanWithProbes } from "../../models/scan";
 import { Probe, ProbeStatus } from "../../models/probe";
 import { Scan } from '../../models/scan'
-import { Report, SupabaseReport } from '../../models/report'
+import { SupabaseReport } from '../../models/report'
 import { publishProbeRequest } from "../../storage/awsSqsQueue";
 import { createReport } from "../../storage/report.storage";
-import { ScanDoesNotExist } from "../../exceptions/exceptions";
+import { ScanDoesNotExist, UserHasNotEnoughCredits } from "../../exceptions/exceptions";
 import { deleteProbes, updateProbesByScanId } from "../../storage/probe.storage";
+import { getUserCredits, updateUserCredits } from "../../storage/credits.storage";
 
 
 export const requestScan = async (scanRequest: CreateScanRequest): Promise<ScanRequestResponse> => {
     const newScanId = uuidv4();
     console.log(`[REQUEST][${newScanId}] Received scan request ${newScanId}`)
 
+    const userCredits = await getUserCredits(scanRequest.user_id)
+
+    if (userCredits < scanRequest.probes?.length) {
+        throw new UserHasNotEnoughCredits(scanRequest.user_id)
+    }
+
     // Assing uids to probes
     const probes: Partial<Probe>[] = scanRequest.probes.map((probe) => {
         return {
@@ -34,7 +41,7 @@ export const requestScan = async (scanRequest: CreateScanRequest): Promise<ScanR
         userId: scanRequest.user_id
     })
     const currentReport = await setupScanNewReport(newScan)
-    await scanStorage.updateScan(newScan.id, { currentReportId: currentReport.id})
+    await scanStorage.updateScan(newScan.id, { currentReportId: currentReport.id })
     console.log(`[REQUEST][${newScanId}] Scan start data saved !`)
 
     console.log(`[REQUEST][${newScanId}] Saving probe start data...`)
@@ -59,6 +66,9 @@ export const requestScan = async (scanRequest: CreateScanRequest): Promise<ScanR
     })));
     console.log(`[REQUEST][${newScanId}] Published request to Queue !`)
 
+    // Update user credits
+    await updateUserCredits(scanRequest.user_id, userCredits - scanRequest.probes?.length)
+
     return { scanId: newScanId };
 }
 
@@ -120,8 +130,13 @@ export const updateScan = async (scanId: string, scanWithProbes: UpdateScanReque
 
 export const restartScan = async (scan: ScanWithProbes): Promise<void> => {
     console.log(`[REQUEST][SCAN][RESTART][${scan.id}] Restarting scan...`)
+    const userCredits = await getUserCredits(scan.userId)
+    if (userCredits < scan.probes?.length) {
+        throw new UserHasNotEnoughCredits(scan.userId)
+    }
+
     const report = await setupScanNewReport(scan)
-    await scanStorage.updateScan(scan.id, { currentReportId: report.id})
+    await scanStorage.updateScan(scan.id, { currentReportId: report.id })
 
     console.log(`[REQUEST][SCAN][RESTART][${scan.id}] Updating probes...`)
     await updateProbesByScanId(scan.id, { status: ProbeStatus.PENDING })
@@ -136,4 +151,7 @@ export const restartScan = async (scan: ScanWithProbes): Promise<void> => {
         settings: probe.settings
     })));
     console.log(`[REQUEST][SCAN][RESTART][${scan.id}] Published request to Queue !`)
+
+    // Update user credits
+    await updateUserCredits(scan.userId, userCredits - scan.probes?.length)
 }

src/services/requests/server.ts

@@ -7,6 +7,7 @@ import { z } from 'zod'
 import { config as loadLocalEnv } from "dotenv";
 import { requireEnvVars } from "../../utils";
 import { isProd } from "../../config";
+import { UserHasNotEnoughCredits } from "../../exceptions/exceptions";
 if (!isProd) {
 	loadLocalEnv();
 }
@@ -49,6 +50,8 @@ server.post("/scans", async (req, res) => {
 		if (e instanceof z.ZodError) {
 			console.warn(e)
 			res.status(400).send(e);
+		} else if (e instanceof UserHasNotEnoughCredits) {
+			res.status(406).send('Not enough credit to run scan')
 		} else {
 			console.error(e)
 			res.status(500).send(e);

src/services/stats/stats.service.ts

@@ -1,17 +1,15 @@
 import { UserStats } from "../../models/stats"
-import { getReportsById } from "../../storage/mongo/mongoReport.storage"
 import { getReportsByUserId } from "../../storage/report.storage"
-import { getScanByUserID } from "../../storage/scan.storage"
+import { getScansWithProbesByUserId } from "../../storage/scan.storage"
 import { isDefined } from "../../utils"
 
 export const getStatsForUser = async (userId: string): Promise<UserStats> => {
     const [scans, reports] = await Promise.all([
-        getScanByUserID(userId),
+        getScansWithProbesByUserId(userId),
         getReportsByUserId(userId)
     ])
 
-    const reportsResults = await getReportsById(reports.map((report) => report.reportId))
-    const nbTotalProbes = reportsResults.reduce((acc, report) => acc + report.nbProbes, 0)
+    const nbTotalProbes = scans.reduce((sum, scan) => sum + scan.probes?.length, 0)
     const avgProbesPerScan = nbTotalProbes / scans.length
 
     return {

src/storage/credits.storage.ts

@@ -0,0 +1,12 @@
+import supabaseClient from "./supabase"
+
+export const getUserCredits = async (userId: string): Promise<number> => {
+    return (await supabaseClient.from('user_credits').select('remaningCredits').eq('userId', userId).maybeSingle()).data?.remaningCredits || 0
+}
+
+export const updateUserCredits = async (userId: string, remaningCredits: number): Promise<void> => {
+    const res = await supabaseClient.from('user_credits').upsert({ userId, remaningCredits }).eq('userId', userId)
+    if (res.error) {
+        console.error(res.error)
+    }
+}

src/storage/scan.storage.ts

@@ -36,8 +36,8 @@ export const listenScans = (onChange: (payload: RealtimePostgresChangesPayload<S
         .subscribe()
 }
 
-export const getScanByUserID = async (userId: string): Promise<Scan[]> => {
-    const res = await supabaseClient.from('scans').select('*').eq('userId', userId)
+export const getScansWithProbesByUserId = async (userId: string): Promise<ScanWithProbes[]> => {
+    const res = await supabaseClient.from('scans').select('*, probes(*)').eq('userId', userId)
     if (res.error) {
         console.log(res.error)
         return []

src/storage/settings.storage.ts

@@ -0,0 +1,9 @@
+import { RealtimePostgresChangesPayload } from "@supabase/supabase-js"
+import supabaseClient from './supabase'
+import { UserSettings } from "../models/settings"
+
+export const listenSettings = (onChange: (payload: RealtimePostgresChangesPayload<UserSettings>) => void) => {
+    supabaseClient.channel('user_settings')
+        .on('postgres_changes', { event: '*', schema: 'public', table: 'user_settings' }, onChange)
+        .subscribe()
+}

supabase/migrations/20230514174039_user_credits.sql

@@ -0,0 +1,23 @@
+create table "public"."user_credits" (
+    "id" bigint generated by default as identity not null,
+    "userId" uuid,
+    "remaningCredits" smallint default '0'::smallint
+);
+
+
+alter table "public"."user_credits" enable row level security;
+
+CREATE UNIQUE INDEX user_credits_pkey ON public.user_credits USING btree (id);
+
+alter table "public"."user_credits" add constraint "user_credits_pkey" PRIMARY KEY using index "user_credits_pkey";
+
+alter table "public"."user_credits" add constraint "user_credits_userId_fkey" FOREIGN KEY ("userId") REFERENCES auth.users(id) ON DELETE CASCADE not valid;
+
+alter table "public"."user_credits" validate constraint "user_credits_userId_fkey";
+
+create policy "Enable select for user"
+on "public"."user_credits"
+as permissive
+for select
+to authenticated
+using ((auth.uid() = "userId"));

supabase/migrations/20230521182054_settings_refactor.sql

@@ -0,0 +1,15 @@
+alter table "public"."user_credits" drop constraint "user_credits_pkey";
+
+drop index if exists "public"."user_credits_pkey";
+
+alter table "public"."user_credits" drop column "id";
+
+alter table "public"."user_credits" alter column "userId" set not null;
+
+CREATE UNIQUE INDEX "user_credits_userId_key" ON public.user_credits USING btree ("userId");
+
+CREATE UNIQUE INDEX user_credits_pkey ON public.user_credits USING btree ("userId");
+
+alter table "public"."user_credits" add constraint "user_credits_pkey" PRIMARY KEY using index "user_credits_pkey";
+
+alter table "public"."user_credits" add constraint "user_credits_userId_key" UNIQUE using index "user_credits_userId_key";