add check for audit with public access and entitlement

Author: msukkari

packages/web/src/app/api/(server)/ee/audit/route.ts

@@ -6,6 +6,8 @@ import { isServiceError } from "@/lib/utils";
 import { serviceErrorResponse } from "@/lib/serviceError";
 import { StatusCodes } from "http-status-codes";
 import { ErrorCode } from "@/lib/errorCodes";
+import { env } from "@/env.mjs";
+import { getEntitlements } from "@/features/entitlements/server";
 
 export const GET = async (request: NextRequest) => {
   const domain = request.headers.get("X-Org-Domain");
@@ -18,6 +20,23 @@ export const GET = async (request: NextRequest) => {
       message: "Missing X-Org-Domain header",
     });
   } 
+
+  if (env.SOURCEBOT_EE_AUDIT_LOGGING_ENABLED === 'false') {
+    return serviceErrorResponse({
+      statusCode: StatusCodes.FORBIDDEN,
+      errorCode: ErrorCode.NOT_FOUND,
+      message: "Audit logging is not enabled",
+    });
+  }
+
+  const entitlements = getEntitlements();
+  if (!entitlements.includes('audit')) {
+    return serviceErrorResponse({
+      statusCode: StatusCodes.FORBIDDEN,
+      errorCode: ErrorCode.NOT_FOUND,
+      message: "Audit logging is not enabled for your license",
+    });
+  }
 
   const result = await fetchAuditRecords(domain, apiKey);
   if (isServiceError(result)) {

packages/web/src/ee/features/audit/factory.ts

@@ -1,16 +1,11 @@
 import { IAuditService } from '@/ee/features/audit/types';
 import { MockAuditService } from '@/ee/features/audit/mockAuditService';
-import { hasEntitlement } from '@/features/entitlements/server';
 import { AuditService } from '@/ee/features/audit/auditService';
+import { env } from '@/env.mjs';
 
-let enterpriseService: IAuditService | null = null;
+let enterpriseService: IAuditService | undefined;
 
 export function getAuditService(): IAuditService {
-  if (hasEntitlement('audit')) {
-    if (!enterpriseService) {
-      enterpriseService = new AuditService();
-    }
-    return enterpriseService!;
-  }
-  return new MockAuditService();
+  enterpriseService = enterpriseService ?? (env.SOURCEBOT_EE_AUDIT_LOGGING_ENABLED === 'true' ? new AuditService() : new MockAuditService());
+  return enterpriseService;
 }

packages/web/src/env.mjs

@@ -84,6 +84,7 @@ export const env = createEnv({
 
         // EE License
         SOURCEBOT_EE_LICENSE_KEY: z.string().optional(),
+        SOURCEBOT_EE_AUDIT_LOGGING_ENABLED: booleanSchema.default('false'),
 
         // GitHub app for review agent
         GITHUB_APP_ID: z.string().optional(),

packages/web/src/initialize.ts

@@ -10,7 +10,7 @@ import { ConnectionConfig } from '@sourcebot/schemas/v3/connection.type';
 import { indexSchema } from '@sourcebot/schemas/v3/index.schema';
 import Ajv from 'ajv';
 import { syncSearchContexts } from '@/ee/features/searchContexts/syncSearchContexts';
-import { hasEntitlement } from '@/features/entitlements/server';
+import { getEntitlements, hasEntitlement } from '@/features/entitlements/server';
 import { createGuestUser, setPublicAccessStatus } from '@/ee/features/publicAccess/publicAccess';
 import { isServiceError } from './lib/utils';
 import { ServiceErrorException } from './lib/serviceError';
@@ -150,6 +150,11 @@ const syncDeclarativeConfig = async (configPath: string) => {
     }
 
     if (hasPublicAccessEntitlement) {
+        if (enablePublicAccess && env.SOURCEBOT_EE_AUDIT_LOGGING_ENABLED === 'true') {
+            logger.error(`Audit logging is not supported when public access is enabled. Please disable audit logging or disable public access.`);
+            process.exit(1);
+        }
+        
         logger.info(`Setting public access status to ${!!enablePublicAccess} for org ${SINGLE_TENANT_ORG_DOMAIN}`);
         const res = await setPublicAccessStatus(SINGLE_TENANT_ORG_DOMAIN, !!enablePublicAccess);
         if (isServiceError(res)) {
@@ -186,6 +191,17 @@ const pruneOldGuestUser = async () => {
     }
 }
 
+const validateEntitlements = () => {
+    const entitlements = getEntitlements();
+
+    if (env.SOURCEBOT_EE_AUDIT_LOGGING_ENABLED === 'true') {
+        if (!hasEntitlement('audit')) {
+            logger.error(`Audit logging is enabled but your license does not include the audit logging entitlement. Please reach out to us to enquire about upgrading your license.`);
+            process.exit(1);
+        }
+    }
+}
+
 const initSingleTenancy = async () => {
     await prisma.org.upsert({
         where: {
@@ -203,6 +219,9 @@ const initSingleTenancy = async () => {
     // To keep things simple, we'll just delete the old guest user if it exists in the DB
     await pruneOldGuestUser();
 
+    // Startup time entitlement/environment variable validation
+    validateEntitlements();
+
     const hasPublicAccessEntitlement = hasEntitlement("public-access");
     if (hasPublicAccessEntitlement) {
         const res = await createGuestUser(SINGLE_TENANT_ORG_DOMAIN);