Merge pull request #645 from Infisical/environment-paywall

Update getPlan to consider the user's current workspace

backend/src/controllers/v1/membershipOrgController.ts

@@ -99,7 +99,7 @@ export const inviteUserToOrganization = async (req: Request, res: Response) => {
     throw new Error('Failed to validate organization membership');
   }
   
-  const plan = await EELicenseService.getOrganizationPlan(organizationId);
+  const plan = await EELicenseService.getPlan(organizationId);
   
   if (plan.memberLimit !== null) {
     // case: limit imposed on number of members allowed

backend/src/controllers/v1/workspaceController.ts

@@ -116,7 +116,7 @@ export const createWorkspace = async (req: Request, res: Response) => {
     throw new Error("Failed to validate organization membership");
   }
 
-  const plan = await EELicenseService.getOrganizationPlan(organizationId);
+  const plan = await EELicenseService.getPlan(organizationId);
   
   if (plan.workspaceLimit !== null) {
     // case: limit imposed on number of workspaces allowed

backend/src/controllers/v2/environmentController.ts

@@ -8,6 +8,7 @@ import {
   Membership,
 } from '../../models';
 import { SecretVersion } from '../../ee/models';
+import { EELicenseService } from '../../ee/services';
 import { BadRequestError } from '../../utils/errors';
 import _ from 'lodash';
 import { PERMISSION_READ_SECRETS, PERMISSION_WRITE_SECRETS } from '../../variables';
@@ -40,6 +41,8 @@ export const createWorkspaceEnvironment = async (
   });
   await workspace.save();
 
+  await EELicenseService.refreshPlan(workspace.organization.toString(), workspaceId);
+
   return res.status(200).send({
     message: 'Successfully created new environment',
     workspace: workspaceId,
@@ -186,7 +189,9 @@ export const deleteWorkspaceEnvironment = async (
   await Membership.updateMany(
     { workspace: workspaceId },
     { $pull: { deniedPermissions: { environmentSlug: environmentSlug } } }
-  )
+  );
+
+  await EELicenseService.refreshPlan(workspace.organization.toString(), workspaceId);
 
   return res.status(200).send({
     message: 'Successfully deleted environment',

backend/src/ee/controllers/v1/organizationsController.ts

@@ -8,8 +8,9 @@ import { EELicenseService } from '../../services';
  */
 export const getOrganizationPlan = async (req: Request, res: Response) => {
     const { organizationId } = req.params;
+    const workspaceId = req.query.workspaceId as string;
 
-    const plan = await EELicenseService.getOrganizationPlan(organizationId);
+    const plan = await EELicenseService.getPlan(organizationId, workspaceId);
 
     return res.status(200).send({
         plan,

backend/src/ee/routes/v1/organizations.ts

@@ -5,7 +5,7 @@ import {
     requireOrganizationAuth,
     validateRequest
 } from '../../../middleware';
-import { param, body } from 'express-validator';
+import { param, body, query } from 'express-validator';
 import { organizationsController } from '../../controllers/v1';
 import {
     OWNER, ADMIN, MEMBER, ACCEPTED
@@ -21,6 +21,7 @@ router.get(
         acceptedStatuses: [ACCEPTED]
     }),
     param('organizationId').exists().trim(),
+    query('workspaceId').optional().isString(),
     validateRequest,
     organizationsController.getOrganizationPlan
 );

backend/src/ee/services/EELicenseService.ts

@@ -22,6 +22,8 @@ interface FeatureSet {
     workspacesUsed: number;
     memberLimit: number | null;
     membersUsed: number;
+    environmentLimit: number | null;
+    environmentsUsed: number;
     secretVersioning: boolean;
     pitRecovery: boolean;
     rbac: boolean;
@@ -51,6 +53,8 @@ class EELicenseService {
         workspacesUsed: 0,
         memberLimit: null,
         membersUsed: 0,
+        environmentLimit: null,
+        environmentsUsed: 0,
         secretVersioning: true,
         pitRecovery: true,
         rbac: true,
@@ -69,10 +73,10 @@ class EELicenseService {
         });
     }
     
-    public async getOrganizationPlan(organizationId: string): Promise<FeatureSet> {
+    public async getPlan(organizationId: string, workspaceId?: string): Promise<FeatureSet> {
         try {
             if (this.instanceType === 'cloud') {
-                const cachedPlan = this.localFeatureSet.get<FeatureSet>(organizationId);
+                const cachedPlan = this.localFeatureSet.get<FeatureSet>(`${organizationId}-${workspaceId ?? ''}`);
                 if (cachedPlan) {
                     return cachedPlan;
                 }
@@ -80,12 +84,16 @@ class EELicenseService {
                 const organization = await Organization.findById(organizationId);
                 if (!organization) throw OrganizationNotFoundError();
 
-                const { data: { currentPlan } } = await licenseServerKeyRequest.get(
-                    `${await getLicenseServerUrl()}/api/license-server/v1/customers/${organization.customerId}/cloud-plan`
-                );
+                let url = `${await getLicenseServerUrl()}/api/license-server/v1/customers/${organization.customerId}/cloud-plan`;
+
+                if (workspaceId) {
+                    url += `?workspaceId=${workspaceId}`;
+                }
+
+                const { data: { currentPlan } } = await licenseServerKeyRequest.get(url);
 
                 // cache fetched plan for organization
-                this.localFeatureSet.set(organizationId, currentPlan);
+                this.localFeatureSet.set(`${organizationId}-${workspaceId ?? ''}`, currentPlan);
 
                 return currentPlan;
             }
@@ -96,10 +104,10 @@ class EELicenseService {
         return this.globalFeatureSet;
     }
     
-    public async refreshOrganizationPlan(organizationId: string) {
+    public async refreshPlan(organizationId: string, workspaceId?: string) {
         if (this.instanceType === 'cloud') {
-            this.localFeatureSet.del(organizationId);
-            await this.getOrganizationPlan(organizationId);
+            this.localFeatureSet.del(`${organizationId}-${workspaceId ?? ''}`);
+            await this.getPlan(organizationId, workspaceId);
         }
     }
 

backend/src/helpers/organization.ts

@@ -170,7 +170,7 @@ export const updateSubscriptionOrgQuantity = async ({
     );
   }
 
-  await EELicenseService.refreshOrganizationPlan(organizationId);
+  await EELicenseService.refreshPlan(organizationId);
 
   return stripeSubscription;
 };

backend/src/helpers/workspace.ts

@@ -41,7 +41,7 @@ export const createWorkspace = async ({
 		workspaceId: workspace._id
 	});
 
-	await EELicenseService.refreshOrganizationPlan(organizationId);
+	await EELicenseService.refreshPlan(organizationId);
 
 	return workspace;
 };