Implement usage tracking on resources, clients

server/lib/billing/features.ts

@@ -5,6 +5,9 @@ export enum LimitId {
     DOMAINS = "domains",
     REMOTE_EXIT_NODES = "remoteExitNodes",
     ORGANIZATIONS = "organizations",
+    PUBLIC_RESOURCES = "publicResources",
+    PRIVATE_RESOURCES = "privateResources",
+    MACHINE_CLIENTS = "machineClients",
     TIER1 = "tier1"
 }
 

server/lib/billing/limitSet.ts

@@ -12,7 +12,10 @@ export const freeLimitSet: LimitSet = {
     [LimitId.USERS]: { value: 5, description: "Basic limit" },
     [LimitId.DOMAINS]: { value: 5, description: "Basic limit" },
     [LimitId.REMOTE_EXIT_NODES]: { value: 1, description: "Basic limit" },
-    [LimitId.ORGANIZATIONS]: { value: 1, description: "Basic limit" }
+    [LimitId.ORGANIZATIONS]: { value: 1, description: "Basic limit" },
+    [LimitId.PUBLIC_RESOURCES]: { value: 15, description: "Basic limit" },
+    [LimitId.PRIVATE_RESOURCES]: { value: 15, description: "Basic limit" },
+    [LimitId.MACHINE_CLIENTS]: { value: 5, description: "Basic limit" }
 };
 
 export const tier1LimitSet: LimitSet = {
@@ -20,7 +23,10 @@ export const tier1LimitSet: LimitSet = {
     [LimitId.SITES]: { value: 10, description: "Home limit" },
     [LimitId.DOMAINS]: { value: 10, description: "Home limit" },
     [LimitId.REMOTE_EXIT_NODES]: { value: 1, description: "Home limit" },
-    [LimitId.ORGANIZATIONS]: { value: 1, description: "Home limit" }
+    [LimitId.ORGANIZATIONS]: { value: 1, description: "Home limit" },
+    [LimitId.PUBLIC_RESOURCES]: { value: 30, description: "Home limit" },
+    [LimitId.PRIVATE_RESOURCES]: { value: 30, description: "Home limit" },
+    [LimitId.MACHINE_CLIENTS]: { value: 10, description: "Home limit" }
 };
 
 export const tier2LimitSet: LimitSet = {
@@ -43,7 +49,10 @@ export const tier2LimitSet: LimitSet = {
     [LimitId.ORGANIZATIONS]: {
         value: 1,
         description: "Team limit"
-    }
+    },
+    [LimitId.PUBLIC_RESOURCES]: { value: 150, description: "Team limit" },
+    [LimitId.PRIVATE_RESOURCES]: { value: 150, description: "Team limit" },
+    [LimitId.MACHINE_CLIENTS]: { value: 25, description: "Team limit" }
 };
 
 export const tier3LimitSet: LimitSet = {
@@ -66,5 +75,8 @@ export const tier3LimitSet: LimitSet = {
     [LimitId.ORGANIZATIONS]: {
         value: 5,
         description: "Business limit"
-    }
+    },
+    [LimitId.PUBLIC_RESOURCES]: { value: 750, description: "Business limit" },
+    [LimitId.PRIVATE_RESOURCES]: { value: 750, description: "Business limit" },
+    [LimitId.MACHINE_CLIENTS]: { value: 100, description: "Business limit" }
 };

server/routers/client/createClient.ts

@@ -30,6 +30,8 @@ import {
 } from "@server/lib/rebuildClientAssociations";
 import { getUniqueClientName } from "@server/db/names";
 import { build } from "@server/build";
+import { LimitId } from "@server/lib/billing";
+import { usageService } from "@server/lib/billing/usageService";
 
 const createClientParamsSchema = z.strictObject({
     orgId: z.string()
@@ -128,6 +130,38 @@ export async function createClient(
             );
         }
 
+        if (build == "saas") {
+            const usage = await usageService.getUsage(
+                orgId,
+                LimitId.MACHINE_CLIENTS
+            );
+            if (!usage) {
+                return next(
+                    createHttpError(
+                        HttpCode.NOT_FOUND,
+                        "No usage data found for this organization"
+                    )
+                );
+            }
+            const rejectClient = await usageService.checkLimitSet(
+                orgId,
+
+                LimitId.MACHINE_CLIENTS,
+                {
+                    ...usage,
+                    instantaneousValue: (usage.instantaneousValue || 0) + 1
+                } // We need to add one to know if we are violating the limit
+            );
+            if (rejectClient) {
+                return next(
+                    createHttpError(
+                        HttpCode.FORBIDDEN,
+                        "Public resource limit exceeded. Please upgrade your plan."
+                    )
+                );
+            }
+        }
+
         const [org] = await db.select().from(orgs).where(eq(orgs.orgId, orgId));
 
         if (!org) {
@@ -289,6 +323,8 @@ export async function createClient(
                 clientId: newClient.clientId,
                 dateCreated: moment().toISOString()
             });
+
+            await usageService.add(orgId, LimitId.MACHINE_CLIENTS, 1, trx);
         });
 
         if (newClient) {
@@ -303,7 +339,7 @@ export async function createClient(
             data: newClient,
             success: true,
             error: false,
-            message: "Site created successfully",
+            message: "Client created successfully",
             status: HttpCode.CREATED
         });
     } catch (error) {

server/routers/client/deleteClient.ts

@@ -15,6 +15,8 @@ import {
 } from "@server/lib/rebuildClientAssociations";
 import { sendTerminateClient } from "./terminate";
 import { OlmErrorCodes } from "../olm/error";
+import { LimitId } from "@server/lib/billing/features";
+import { usageService } from "@server/lib/billing/usageService";
 
 const deleteClientSchema = z.strictObject({
     clientId: z.coerce.number().int().positive()
@@ -118,6 +120,13 @@ export async function deleteClient(
             if (!client.userId && client.olmId) {
                 await trx.delete(olms).where(eq(olms.olmId, client.olmId));
             }
+
+            await usageService.add(
+                deletedClient.orgId,
+                LimitId.MACHINE_CLIENTS,
+                -1,
+                trx
+            );
         });
 
         if (deletedClient) {

server/routers/resource/createResource.ts

@@ -36,6 +36,8 @@ import {
     getUniqueResourceName,
     getUniqueResourcePolicyName
 } from "@server/db/names";
+import { usageService } from "@server/lib/billing/usageService";
+import { LimitId } from "@server/lib/billing";
 
 const createResourceParamsSchema = z.strictObject({
     orgId: z.string()
@@ -235,6 +237,38 @@ export async function createResource(
             req.body.mode = resolvedMode.mode;
         }
 
+        if (build == "saas") {
+            const usage = await usageService.getUsage(
+                orgId,
+                LimitId.PUBLIC_RESOURCES
+            );
+            if (!usage) {
+                return next(
+                    createHttpError(
+                        HttpCode.NOT_FOUND,
+                        "No usage data found for this organization"
+                    )
+                );
+            }
+            const rejectResource = await usageService.checkLimitSet(
+                orgId,
+
+                LimitId.PUBLIC_RESOURCES,
+                {
+                    ...usage,
+                    instantaneousValue: (usage.instantaneousValue || 0) + 1
+                } // We need to add one to know if we are violating the limit
+            );
+            if (rejectResource) {
+                return next(
+                    createHttpError(
+                        HttpCode.FORBIDDEN,
+                        "Public resource limit exceeded. Please upgrade your plan."
+                    )
+                );
+            }
+        }
+
         if (typeof req.body.proxyPort === "number") {
             if (
                 !config.getRawConfig().flags?.allow_raw_resources &&
@@ -503,6 +537,8 @@ async function createHttpResource(
         }
 
         resource = newResource[0];
+
+        await usageService.add(orgId, LimitId.PUBLIC_RESOURCES, 1, trx);
     });
 
     if (!resource) {
@@ -631,6 +667,8 @@ async function createRawResource(
         }
 
         resource = newResource[0];
+
+        await usageService.add(orgId, LimitId.PUBLIC_RESOURCES, 1, trx);
     });
 
     if (!resource) {

server/routers/resource/deleteResource.ts

@@ -11,6 +11,8 @@ import {
     performDeleteResource,
     runResourceDeleteSideEffects
 } from "@server/lib/deleteResource";
+import { LimitId } from "@server/lib/billing";
+import { usageService } from "@server/lib/billing/usageService";
 
 const deleteResourceSchema = z.strictObject({
     resourceId: z.coerce.number().int().positive()
@@ -64,6 +66,14 @@ export async function deleteResource(
 
         await db.transaction(async (trx) => {
             deleteResult = await performDeleteResource(resourceId, trx);
+            if (deleteResult?.deletedResource?.orgId) {
+                await usageService.add(
+                    deleteResult?.deletedResource?.orgId,
+                    LimitId.PUBLIC_RESOURCES,
+                    -1,
+                    trx
+                );
+            }
         });
 
         if (!deleteResult) {

server/routers/siteResource/createSiteResource.ts

@@ -37,6 +37,8 @@ import { fromError } from "zod-validation-error";
 import { validateAndConstructDomain } from "@server/lib/domainUtils";
 import { createCertificate } from "#dynamic/routers/certificates/createCertificate";
 import { build } from "@server/build";
+import { usageService } from "@server/lib/billing/usageService";
+import { LimitId } from "@server/lib/billing";
 
 const createSiteResourceParamsSchema = z.strictObject({
     orgId: z.string()
@@ -294,6 +296,38 @@ export async function createSiteResource(
             siteIds.push(siteId);
         }
 
+        if (build == "saas") {
+            const usage = await usageService.getUsage(
+                orgId,
+                LimitId.PRIVATE_RESOURCES
+            );
+            if (!usage) {
+                return next(
+                    createHttpError(
+                        HttpCode.NOT_FOUND,
+                        "No usage data found for this organization"
+                    )
+                );
+            }
+            const rejectResource = await usageService.checkLimitSet(
+                orgId,
+
+                LimitId.PRIVATE_RESOURCES,
+                {
+                    ...usage,
+                    instantaneousValue: (usage.instantaneousValue || 0) + 1
+                } // We need to add one to know if we are violating the limit
+            );
+            if (rejectResource) {
+                return next(
+                    createHttpError(
+                        HttpCode.FORBIDDEN,
+                        "Private resource limit exceeded. Please upgrade your plan."
+                    )
+                );
+            }
+        }
+
         if (mode == "http") {
             const hasHttpFeature = await isLicensedOrSubscribed(
                 orgId,
@@ -605,6 +639,13 @@ export async function createSiteResource(
                         );
                     }
                 }
+
+                await usageService.add(
+                    orgId,
+                    LimitId.PRIVATE_RESOURCES,
+                    1,
+                    trx
+                );
             });
         } finally {
             await releaseAliasLock?.();

server/routers/siteResource/deleteSiteResource.ts

@@ -12,6 +12,8 @@ import {
     performDeleteSiteResource,
     runSiteResourceDeleteSideEffects
 } from "@server/lib/deleteSiteResource";
+import { LimitId } from "@server/lib/billing";
+import { usageService } from "@server/lib/billing/usageService";
 
 const deleteSiteResourceParamsSchema = z.strictObject({
     siteResourceId: z.coerce.number().int().positive()
@@ -86,6 +88,14 @@ export async function deleteSiteResource(
                 siteResourceId,
                 trx
             );
+            if (removedSiteResource?.orgId) {
+                await usageService.add(
+                    removedSiteResource?.orgId,
+                    LimitId.PRIVATE_RESOURCES,
+                    -1,
+                    trx
+                );
+            }
         });
 
         if (!removedSiteResource) {