Merge pull request #3057 from trycompai/main

[comp] Production Deploy

Author: tofikwest

apps/api/src/auth/service-token-only.guard.spec.ts

@@ -0,0 +1,45 @@
+import { ExecutionContext, ForbiddenException } from '@nestjs/common';
+import { ServiceTokenOnlyGuard } from './service-token-only.guard';
+import type { AuthenticatedRequest } from './types';
+
+function contextFor(
+  request: Partial<AuthenticatedRequest>,
+): ExecutionContext {
+  return {
+    switchToHttp: () => ({
+      getRequest: () => request as AuthenticatedRequest,
+    }),
+  } as unknown as ExecutionContext;
+}
+
+describe('ServiceTokenOnlyGuard', () => {
+  const guard = new ServiceTokenOnlyGuard();
+
+  it('allows internal service-token requests', () => {
+    expect(
+      guard.canActivate(contextFor({ isServiceToken: true })),
+    ).toBe(true);
+  });
+
+  it('rejects session requests', () => {
+    expect(() =>
+      guard.canActivate(
+        contextFor({ authType: 'session', isApiKey: false, isServiceToken: false }),
+      ),
+    ).toThrow(ForbiddenException);
+  });
+
+  it('rejects API-key requests', () => {
+    expect(() =>
+      guard.canActivate(
+        contextFor({ authType: 'api-key', isApiKey: true, isServiceToken: false }),
+      ),
+    ).toThrow(ForbiddenException);
+  });
+
+  it('rejects requests with no service-token flag set', () => {
+    expect(() => guard.canActivate(contextFor({}))).toThrow(
+      ForbiddenException,
+    );
+  });
+});

apps/api/src/auth/service-token-only.guard.ts

@@ -0,0 +1,30 @@
+import {
+  CanActivate,
+  ExecutionContext,
+  ForbiddenException,
+  Injectable,
+} from '@nestjs/common';
+import { AuthenticatedRequest } from './types';
+
+/**
+ * Guard that rejects everything except internal service-token auth.
+ * Use on internal-only endpoints that return sensitive data not meant for
+ * customers (e.g., minting temporary AWS credentials for background workers).
+ *
+ * Place between HybridAuthGuard and PermissionGuard:
+ * @UseGuards(HybridAuthGuard, ServiceTokenOnlyGuard, PermissionGuard)
+ */
+@Injectable()
+export class ServiceTokenOnlyGuard implements CanActivate {
+  canActivate(context: ExecutionContext): boolean {
+    const request = context.switchToHttp().getRequest<AuthenticatedRequest>();
+
+    if (!request.isServiceToken) {
+      throw new ForbiddenException(
+        'This endpoint is only available for internal service-token requests.',
+      );
+    }
+
+    return true;
+  }
+}

apps/api/src/cloud-security/cloud-security.controller.spec.ts

@@ -16,7 +16,10 @@ import { Test, TestingModule } from '@nestjs/testing';
 import { HttpException, HttpStatus } from '@nestjs/common';
 
 import { CloudSecurityController } from './cloud-security.controller';
-import { CloudSecurityService } from './cloud-security.service';
+import {
+  CloudSecurityService,
+  ConnectionNotFoundError,
+} from './cloud-security.service';
 import { CloudSecurityQueryService } from './cloud-security-query.service';
 import { CloudSecurityLegacyService } from './cloud-security-legacy.service';
 import { CloudSecurityActivityService } from './cloud-security-activity.service';
@@ -29,6 +32,7 @@ import { CloudAwsScanModeService } from './aws-scan-mode.service';
 import { ActingUserResolver } from '../auth/acting-user.service';
 import { HybridAuthGuard } from '../auth/hybrid-auth.guard';
 import { PermissionGuard } from '../auth/permission.guard';
+import { ServiceTokenOnlyGuard } from '../auth/service-token-only.guard';
 import type { AuthenticatedRequest } from '../auth/types';
 
 /**
@@ -87,6 +91,8 @@ describe('CloudSecurityController — API-key mutation support', () => {
       .useValue(mockGuard)
       .overrideGuard(PermissionGuard)
       .useValue(mockGuard)
+      .overrideGuard(ServiceTokenOnlyGuard)
+      .useValue(mockGuard)
       .compile();
 
     controller = module.get(CloudSecurityController);
@@ -309,3 +315,88 @@ describe('CloudSecurityController — API-key mutation support', () => {
     });
   });
 });
+
+describe('CloudSecurityController — resolveSession (internal)', () => {
+  let controller: CloudSecurityController;
+  let cloudSecurityService: { resolveAwsSession: jest.Mock };
+  const mockGuard = { canActivate: jest.fn().mockReturnValue(true) };
+
+  beforeEach(async () => {
+    cloudSecurityService = { resolveAwsSession: jest.fn() };
+
+    const module: TestingModule = await Test.createTestingModule({
+      controllers: [CloudSecurityController],
+      providers: [
+        { provide: CloudSecurityService, useValue: cloudSecurityService },
+        { provide: CloudSecurityQueryService, useValue: {} },
+        { provide: CloudSecurityLegacyService, useValue: {} },
+        { provide: CloudSecurityActivityService, useValue: {} },
+        { provide: GCPSecurityService, useValue: {} },
+        { provide: AzureSecurityService, useValue: {} },
+        { provide: CheckDefinitionService, useValue: {} },
+        { provide: CloudExceptionService, useValue: {} },
+        { provide: CloudHistoryService, useValue: {} },
+        { provide: CloudAwsScanModeService, useValue: {} },
+        { provide: ActingUserResolver, useValue: {} },
+      ],
+    })
+      .overrideGuard(HybridAuthGuard)
+      .useValue(mockGuard)
+      .overrideGuard(PermissionGuard)
+      .useValue(mockGuard)
+      .overrideGuard(ServiceTokenOnlyGuard)
+      .useValue(mockGuard)
+      .compile();
+
+    controller = module.get(CloudSecurityController);
+    jest.clearAllMocks();
+  });
+
+  it('returns the resolved session from the service', async () => {
+    const session = {
+      ok: true,
+      session: {
+        accessKeyId: 'AKIA_TEMP',
+        secretAccessKey: 'secret',
+        sessionToken: 'token',
+      },
+    };
+    cloudSecurityService.resolveAwsSession.mockResolvedValueOnce(session);
+
+    const result = await controller.resolveSession('conn_1', 'org_1');
+
+    expect(cloudSecurityService.resolveAwsSession).toHaveBeenCalledWith(
+      'conn_1',
+      'org_1',
+    );
+    expect(result).toEqual(session);
+  });
+
+  it('passes through not_configured / assume_failed results', async () => {
+    cloudSecurityService.resolveAwsSession.mockResolvedValueOnce({
+      ok: false,
+      reason: 'assume_failed',
+      error: 'denied',
+    });
+
+    const result = await controller.resolveSession('conn_1', 'org_1');
+    expect(result).toEqual({
+      ok: false,
+      reason: 'assume_failed',
+      error: 'denied',
+    });
+  });
+
+  it('maps ConnectionNotFoundError to a 404', async () => {
+    cloudSecurityService.resolveAwsSession.mockRejectedValueOnce(
+      new ConnectionNotFoundError(),
+    );
+
+    const error = await controller
+      .resolveSession('conn_missing', 'org_1')
+      .catch((e) => e);
+
+    expect(error).toBeInstanceOf(HttpException);
+    expect((error as HttpException).getStatus()).toBe(HttpStatus.NOT_FOUND);
+  });
+});

apps/api/src/cloud-security/cloud-security.controller.ts

@@ -17,6 +17,7 @@ import { ApiOperation } from '@nestjs/swagger';
 import { SkipThrottle } from '@nestjs/throttler';
 import { HybridAuthGuard } from '../auth/hybrid-auth.guard';
 import { PermissionGuard } from '../auth/permission.guard';
+import { ServiceTokenOnlyGuard } from '../auth/service-token-only.guard';
 import { RequirePermission } from '../auth/require-permission.decorator';
 import { OrganizationId } from '../auth/auth-context.decorator';
 import {
@@ -254,6 +255,31 @@ export class CloudSecurityController {
     return { data: definition };
   }
 
+  @Post('resolve-session/:connectionId')
+  @SkipThrottle()
+  @UseGuards(HybridAuthGuard, ServiceTokenOnlyGuard, PermissionGuard)
+  @RequirePermission('integration', 'update')
+  @ApiOperation({
+    summary:
+      'Resolve short-lived AWS credentials for a connection (internal only)',
+  })
+  async resolveSession(
+    @Param('connectionId') connectionId: string,
+    @OrganizationId() organizationId: string,
+  ) {
+    try {
+      return await this.cloudSecurityService.resolveAwsSession(
+        connectionId,
+        organizationId,
+      );
+    } catch (error) {
+      if (error instanceof ConnectionNotFoundError) {
+        throw new HttpException('Connection not found', HttpStatus.NOT_FOUND);
+      }
+      throw error;
+    }
+  }
+
   @Post('scan/:connectionId')
   @UseGuards(HybridAuthGuard, PermissionGuard)
   @RequirePermission('integration', 'update')

apps/api/src/cloud-security/cloud-security.service.resolve-session.spec.ts

@@ -0,0 +1,125 @@
+jest.mock('@db', () => ({
+  db: { integrationConnection: { findFirst: jest.fn() } },
+  Prisma: {},
+}));
+
+import { db } from '@db';
+import {
+  CloudSecurityService,
+  ConnectionNotFoundError,
+} from './cloud-security.service';
+
+type Ctor = ConstructorParameters<typeof CloudSecurityService>;
+
+const findFirst = (db as unknown as {
+  integrationConnection: { findFirst: jest.Mock };
+}).integrationConnection.findFirst;
+
+describe('CloudSecurityService.resolveAwsSession', () => {
+  let credentialVault: { getDecryptedCredentials: jest.Mock };
+  let awsService: { resolveRoleSession: jest.Mock };
+  let service: CloudSecurityService;
+
+  beforeEach(() => {
+    credentialVault = { getDecryptedCredentials: jest.fn() };
+    awsService = { resolveRoleSession: jest.fn() };
+    service = new CloudSecurityService(
+      credentialVault as unknown as Ctor[0],
+      {} as unknown as Ctor[1],
+      {} as unknown as Ctor[2],
+      awsService as unknown as Ctor[3],
+      {} as unknown as Ctor[4],
+      {} as unknown as Ctor[5],
+    );
+    jest.clearAllMocks();
+  });
+
+  it('scopes the connection lookup by organizationId (tenant isolation)', async () => {
+    findFirst.mockResolvedValueOnce(null);
+
+    await expect(
+      service.resolveAwsSession('conn_1', 'org_1'),
+    ).rejects.toBeInstanceOf(ConnectionNotFoundError);
+
+    expect(findFirst).toHaveBeenCalledWith(
+      expect.objectContaining({
+        where: { id: 'conn_1', organizationId: 'org_1', status: 'active' },
+      }),
+    );
+  });
+
+  it('returns not_configured for a non-AWS provider (never assumes)', async () => {
+    findFirst.mockResolvedValueOnce({ provider: { slug: 'gcp' } });
+
+    const result = await service.resolveAwsSession('conn_1', 'org_1');
+
+    expect(result).toEqual({ ok: false, reason: 'not_configured' });
+    expect(awsService.resolveRoleSession).not.toHaveBeenCalled();
+  });
+
+  it('returns not_configured when roleArn/externalId are missing', async () => {
+    findFirst.mockResolvedValueOnce({
+      provider: { slug: 'aws' },
+      variables: {},
+    });
+    credentialVault.getDecryptedCredentials.mockResolvedValueOnce({
+      externalId: 'eid', // roleArn missing
+    });
+
+    const result = await service.resolveAwsSession('conn_1', 'org_1');
+
+    expect(result).toEqual({ ok: false, reason: 'not_configured' });
+    expect(awsService.resolveRoleSession).not.toHaveBeenCalled();
+  });
+
+  it('returns the resolved session on success', async () => {
+    findFirst.mockResolvedValueOnce({
+      provider: { slug: 'aws' },
+      variables: { regions: ['us-east-1'] },
+    });
+    credentialVault.getDecryptedCredentials.mockResolvedValueOnce({
+      roleArn: 'arn:aws:iam::123456789012:role/x',
+      externalId: 'eid',
+      regions: ['us-east-1'],
+    });
+    awsService.resolveRoleSession.mockResolvedValueOnce({
+      accessKeyId: 'AKIA_TEMP',
+      secretAccessKey: 'secret',
+      sessionToken: 'token',
+    });
+
+    const result = await service.resolveAwsSession('conn_1', 'org_1');
+
+    expect(result).toEqual({
+      ok: true,
+      session: {
+        accessKeyId: 'AKIA_TEMP',
+        secretAccessKey: 'secret',
+        sessionToken: 'token',
+      },
+    });
+  });
+
+  it('returns assume_failed with the real reason when the assume throws', async () => {
+    findFirst.mockResolvedValueOnce({
+      provider: { slug: 'aws' },
+      variables: {},
+    });
+    credentialVault.getDecryptedCredentials.mockResolvedValueOnce({
+      roleArn: 'arn:aws:iam::123456789012:role/x',
+      externalId: 'eid',
+      regions: ['us-east-1'],
+    });
+    awsService.resolveRoleSession.mockRejectedValueOnce(
+      new Error('not authorized to perform sts:AssumeRole'),
+    );
+
+    const result = await service.resolveAwsSession('conn_1', 'org_1');
+
+    expect(result).toEqual({
+      ok: false,
+      reason: 'assume_failed',
+      error: 'not authorized to perform sts:AssumeRole',
+    });
+  });
+});