From 96607153dccafacb1d465dd3375fba19d778ec5c Mon Sep 17 00:00:00 2001
From: Tuan Dang <dangtony98@gmail.com>
Date: Tue, 23 May 2023 23:54:54 +0300
Subject: [PATCH] Modularize getOrganizationPlan function

---
 .../controllers/v1/organizationsController.ts | 27 +++----------------
 backend/src/ee/services/EELicenseService.ts   | 25 +++++++++++++++++
 2 files changed, 29 insertions(+), 23 deletions(-)

diff --git a/backend/src/ee/controllers/v1/organizationsController.ts b/backend/src/ee/controllers/v1/organizationsController.ts
index b79379f29..0bdf6e435 100644
--- a/backend/src/ee/controllers/v1/organizationsController.ts
+++ b/backend/src/ee/controllers/v1/organizationsController.ts
@@ -1,4 +1,3 @@
-import * as Sentry from '@sentry/node';
 import { Request, Response } from 'express';
 import { getLicenseServerUrl } from '../../../config';
 import { licenseServerKeyRequest } from '../../../config/request';
@@ -8,31 +7,13 @@ import { EELicenseService } from '../../services';
  * Return the organization's current plan and allowed feature set
  */
 export const getOrganizationPlan = async (req: Request, res: Response) => {
-    try {
-        if (EELicenseService.instanceType === 'cloud') {
-            // instance of Infisical is a cloud instance
+    const plan = await EELicenseService.getOrganizationPlan(req.organization._id.toString());
 
-            const cachedPlan = EELicenseService.localFeatureSet.get(req.organization._id.toString());
-            if (cachedPlan) return cachedPlan;
-
-            const { data: { currentPlan } } = await licenseServerKeyRequest.get(
-                `${await getLicenseServerUrl()}/api/license-server/v1/customers/${req.organization.customerId}/cloud-plan`
-            );
-
-            // cache fetched plan for organization
-            EELicenseService.localFeatureSet.set(req.organization._id.toString(), currentPlan);
-
-            return res.status(200).send({
-                plan: currentPlan
-            });
-        }
-    } catch (err) {
-        Sentry.setUser({ email: req.user.email });
-		Sentry.captureException(err);
-    }
+    // cache fetched plan for organization
+    EELicenseService.localFeatureSet.set(req.organization._id.toString(), plan);
 
     return res.status(200).send({
-        plan: EELicenseService.globalFeatureSet
+        plan
     });
 }
 
diff --git a/backend/src/ee/services/EELicenseService.ts b/backend/src/ee/services/EELicenseService.ts
index 02cab7e5d..7c12931f0 100644
--- a/backend/src/ee/services/EELicenseService.ts
+++ b/backend/src/ee/services/EELicenseService.ts
@@ -7,9 +7,12 @@ import {
 } from '../../config';
 import { 
     licenseKeyRequest,
+    licenseServerKeyRequest,
     refreshLicenseServerKeyToken,
     refreshLicenseKeyToken
 } from '../../config/request';
+import { Organization } from '../../models';
+import { OrganizationNotFoundError } from '../../utils/errors';
 
 interface FeatureSet {
     _id: string | null;
@@ -59,6 +62,28 @@ class EELicenseService {
             stdTTL: 300
         });
     }
+    
+    public async getOrganizationPlan(organizationId: string) {
+        try {
+            if (this.instanceType === 'cloud') {
+                const cachedPlan = this.localFeatureSet.get(organizationId);
+                if (cachedPlan) return cachedPlan;
+
+                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`
+                );
+
+                return currentPlan;
+            }
+        } catch (err) {
+            return this.globalFeatureSet;
+        }
+
+        return this.globalFeatureSet;
+    }
 
     public async initGlobalFeatureSet() {
         const licenseServerKey = await getLicenseServerKey();