consolidate restriction logic inside module

RocketChat · Oct 7, 2024 · 5c7f7ad · 5c7f7ad
1 parent 787eb3e
commit 5c7f7ad
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 45 deletions.
diff --git a/apps/meteor/ee/app/license/server/airGappedRestrictionsCheck.ts b/apps/meteor/ee/app/license/server/airGappedRestrictionsCheck.ts
@@ -1,16 +1,8 @@
-import { License, AirGappedRestriction } from '@rocket.chat/license';
+import { AirGappedRestriction } from '@rocket.chat/license';
 import { Statistics } from '@rocket.chat/models';
 
 export async function checkAirGappedRestrictions(): Promise<void> {
-	if (License.hasModule('unlimited-presence')) {
-		AirGappedRestriction.removeRestrictions();
-		return;
-	}
-
 	const { statsToken: encryptedStatsToken } = (await Statistics.findLast()) || {};
-	if (!encryptedStatsToken) {
-		return AirGappedRestriction.applyRestrictions();
-	}
 
-	await AirGappedRestriction.checkRemainingDaysSinceLastStatsReport(encryptedStatsToken);
+	await AirGappedRestriction.computeRestriction(encryptedStatsToken);
 }
diff --git a/ee/packages/license/src/AirGappedRestriction.spec.ts b/ee/packages/license/src/AirGappedRestriction.spec.ts
@@ -1,13 +1,32 @@
 import { AirGappedRestriction } from './AirGappedRestriction';
 import { StatsTokenBuilder } from './MockedLicenseBuilder';
+import { License } from './licenseImp';
+
+jest.mock('./licenseImp', () => ({
+	License: {
+		hasModule: jest.fn().mockReturnValue(false),
+	},
+}));
 
 describe('AirGappedRestriction', () => {
-	describe('#checkRemainingDaysSinceLastStatsReport()', () => {
+	describe('#computeRestriction()', () => {
+		it('should notify remaining days = 0 (apply restriction) when token is not a string', async () => {
+			const handler = jest.fn();
+
+			AirGappedRestriction.on('remainingDays', handler);
+			await AirGappedRestriction.computeRestriction();
+
+			expect(handler).toHaveBeenCalledTimes(1);
+			expect(handler).toHaveBeenCalledWith({
+				days: 0,
+			});
+			expect(AirGappedRestriction.isRestricted).toBe(true);
+		});
 		it('should notify remaining days = 0 (apply restriction) when it was not possible to decrypt the stats token', async () => {
 			const handler = jest.fn();
 
 			AirGappedRestriction.on('remainingDays', handler);
-			await AirGappedRestriction.checkRemainingDaysSinceLastStatsReport('invalid-token');
+			await AirGappedRestriction.computeRestriction('invalid-token');
 
 			expect(handler).toHaveBeenCalledTimes(1);
 			expect(handler).toHaveBeenCalledWith({
@@ -22,7 +41,7 @@ describe('AirGappedRestriction', () => {
 			const handler = jest.fn();
 
 			AirGappedRestriction.on('remainingDays', handler);
-			await AirGappedRestriction.checkRemainingDaysSinceLastStatsReport(token);
+			await AirGappedRestriction.computeRestriction(token);
 
 			expect(handler).toHaveBeenCalledTimes(1);
 			expect(handler).toHaveBeenCalledWith({
@@ -37,7 +56,7 @@ describe('AirGappedRestriction', () => {
 			const handler = jest.fn();
 
 			AirGappedRestriction.on('remainingDays', handler);
-			await AirGappedRestriction.checkRemainingDaysSinceLastStatsReport(token);
+			await AirGappedRestriction.computeRestriction(token);
 
 			expect(handler).toHaveBeenCalledTimes(1);
 			expect(handler).toHaveBeenCalledWith({
@@ -52,35 +71,20 @@ describe('AirGappedRestriction', () => {
 			const handler = jest.fn();
 
 			AirGappedRestriction.on('remainingDays', handler);
-			await AirGappedRestriction.checkRemainingDaysSinceLastStatsReport(token);
+			await AirGappedRestriction.computeRestriction(token);
 
 			expect(handler).toHaveBeenCalledTimes(1);
 			expect(handler).toHaveBeenCalledWith({
 				days: 0,
 			});
 			expect(AirGappedRestriction.isRestricted).toBe(true);
 		});
-	});
-	describe('#applyRestrictions()', () => {
-		it('should notify remaining days = 0 when restrictions should be applied', () => {
-			const handler = jest.fn();
-
-			AirGappedRestriction.on('remainingDays', handler);
-			AirGappedRestriction.applyRestrictions();
-
-			expect(handler).toHaveBeenCalledTimes(1);
-			expect(handler).toHaveBeenCalledWith({
-				days: 0,
-			});
-			expect(AirGappedRestriction.isRestricted).toBe(true);
-		});
-	});
-	describe('#removeRestrictions()', () => {
-		it('should notify remaining days = -1 when restrictions should be reverted', () => {
+		it('should notify remaining days = -1 when unlimited-presence module is available', () => {
 			const handler = jest.fn();
+			(License.hasModule as jest.Mock).mockReturnValueOnce(true);
 
 			AirGappedRestriction.on('remainingDays', handler);
-			AirGappedRestriction.removeRestrictions();
+			AirGappedRestriction.computeRestriction();
 
 			expect(handler).toHaveBeenCalledTimes(1);
 			expect(handler).toHaveBeenCalledWith({

diff --git a/ee/packages/license/src/AirGappedRestriction.ts b/ee/packages/license/src/AirGappedRestriction.ts
@@ -1,5 +1,6 @@
 import EventEmitter from 'events';
 
+import { License } from '.';
 import { decryptStatsToken } from './token';
 
 const DAY_IN_MS = 24 * 60 * 60 * 1000;
@@ -13,7 +14,21 @@ class AirGappedRestrictionClass extends EventEmitter {
 		return this.restricted;
 	}
 
-	public async checkRemainingDaysSinceLastStatsReport(encryptedToken: string): Promise<void> {
+	public async computeRestriction(encryptedToken?: string): Promise<void> {
+		if (License.hasModule('unlimited-presence')) {
+			this.removeRestrictionsUnderLicense();
+			return;
+		}
+
+		if (typeof encryptedToken !== 'string') {
+			this.applyRestrictions();
+			return;
+		}
+
+		return this.checkRemainingDaysSinceLastStatsReport(encryptedToken);
+	}
+
+	private async checkRemainingDaysSinceLastStatsReport(encryptedToken: string): Promise<void> {
 		try {
 			const { timestamp: lastStatsReportTimestamp } = JSON.parse(await decryptStatsToken(encryptedToken));
 			const now = new Date();
@@ -29,14 +44,13 @@ class AirGappedRestrictionClass extends EventEmitter {
 		}
 	}
 
-	public applyRestrictions(): void {
-		this.restricted = true;
-		this.emit('remainingDays', { days: 0 });
+	private applyRestrictions(): void {
+		this.notifyRemainingDaysUntilRestriction(NO_ACTION_PERIOD_IN_DAYS + WARNING_PERIOD_IN_DAYS);
 	}
 
-	public removeRestrictions(days = -1): void {
+	private removeRestrictionsUnderLicense(): void {
 		this.restricted = false;
-		this.emit('remainingDays', { days });
+		this.emit('remainingDays', { days: -1 });
 	}
 
 	public isWarningPeriod(days: number) {
@@ -47,13 +61,10 @@ class AirGappedRestrictionClass extends EventEmitter {
 	}
 
 	private notifyRemainingDaysUntilRestriction(daysSinceLastStatsReport: number): void {
-		const remainingDaysUntilRestriction = NO_ACTION_PERIOD_IN_DAYS + WARNING_PERIOD_IN_DAYS - daysSinceLastStatsReport;
-		const olderThanTenDays = remainingDaysUntilRestriction <= 0;
-		if (olderThanTenDays) {
-			return this.applyRestrictions();
-		}
+		const remainingDaysUntilRestriction = Math.max(NO_ACTION_PERIOD_IN_DAYS + WARNING_PERIOD_IN_DAYS - daysSinceLastStatsReport, 0);
 
-		this.removeRestrictions(remainingDaysUntilRestriction);
+		this.restricted = remainingDaysUntilRestriction === 0;
+		this.emit('remainingDays', { days: remainingDaysUntilRestriction });
 	}
 }