From 35623ad987bb0647cffe057eb1c6d9ef84e4b3ae Mon Sep 17 00:00:00 2001 From: dal Date: Wed, 13 Aug 2025 14:25:41 -0600 Subject: [PATCH] Integrate access controls with metric file downloads - Added permission checks to metric file download handler - Added permission checks to export-metric-data trigger task - Created comprehensive tests for access control integration - Updated AssetType enums to include report_file - Resolved export conflicts in server-shared package - Fixed AI package imports for new getPermissionedDatasets API --- .../metric_files/download-metric-file.test.ts | 356 ++++++++++++++++++ .../v2/metric_files/download-metric-file.ts | 25 +- .../export-metric-data/export-metric-data.ts | 41 +- .../src/types/asset-permissions.ts | 2 + .../investigation-instructions.ts | 6 +- .../check-asset-permission.ts | 9 +- packages/database/src/schema-types/enums.ts | 1 + packages/server-shared/src/index.ts | 2 + packages/server-shared/src/share/assets.ts | 5 +- .../src/share/individual-permissions.ts | 4 +- packages/server-shared/src/share/requests.ts | 8 +- 11 files changed, 440 insertions(+), 19 deletions(-) create mode 100644 apps/server/src/api/v2/metric_files/download-metric-file.test.ts diff --git a/apps/server/src/api/v2/metric_files/download-metric-file.test.ts b/apps/server/src/api/v2/metric_files/download-metric-file.test.ts new file mode 100644 index 000000000..7d52da301 --- /dev/null +++ b/apps/server/src/api/v2/metric_files/download-metric-file.test.ts @@ -0,0 +1,356 @@ +import { checkPermission } from '@buster/access-controls'; +import { getUserOrganizationId } from '@buster/database'; +import { tasks } from '@trigger.dev/sdk'; +import { HTTPException } from 'hono/http-exception'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { downloadMetricFileHandler } from './download-metric-file'; + +// Mock all external dependencies +vi.mock('@buster/access-controls'); +vi.mock('@buster/database'); +vi.mock('@trigger.dev/sdk', () => ({ + tasks: { + trigger: vi.fn(), + }, + runs: { + retrieve: vi.fn(), + }, +})); + +describe('downloadMetricFileHandler', () => { + const mockUser = { + id: 'user-123', + email: 'test@example.com', + name: 'Test User', + }; + + const mockMetricId = 'metric-456'; + const mockOrganizationId = 'org-789'; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + describe('Access Control', () => { + it('should throw 403 if user is not part of an organization', async () => { + vi.mocked(getUserOrganizationId).mockResolvedValue(null); + + await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toThrow( + HTTPException + ); + + await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toMatchObject({ + status: 403, + message: 'You must be part of an organization to download metric files', + }); + }); + + it('should throw 403 if user does not have permission to view the metric file', async () => { + vi.mocked(getUserOrganizationId).mockResolvedValue({ + organizationId: mockOrganizationId, + userId: mockUser.id, + role: 'member', + } as any); + + vi.mocked(checkPermission).mockResolvedValue({ + hasAccess: false, + effectiveRole: undefined, + accessPath: undefined, + }); + + await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toThrow( + HTTPException + ); + + await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toMatchObject({ + status: 403, + message: 'You do not have permission to download this metric file', + }); + + // Verify permission check was called with correct parameters + expect(checkPermission).toHaveBeenCalledWith({ + userId: mockUser.id, + assetId: mockMetricId, + assetType: 'metric_file', + requiredRole: 'can_view', + organizationId: mockOrganizationId, + }); + }); + + it('should proceed with export when user has permission', async () => { + const mockTaskHandle = { id: 'task-handle-123' }; + + vi.mocked(getUserOrganizationId).mockResolvedValue({ + organizationId: mockOrganizationId, + userId: mockUser.id, + role: 'member', + } as any); + + vi.mocked(checkPermission).mockResolvedValue({ + hasAccess: true, + effectiveRole: 'can_view', + accessPath: 'direct', + }); + + vi.mocked(tasks.trigger).mockResolvedValue(mockTaskHandle as any); + + // Mock successful task completion + const { runs } = await import('@trigger.dev/sdk'); + vi.mocked(runs.retrieve).mockResolvedValue({ + status: 'COMPLETED', + output: { + success: true, + downloadUrl: 'https://example.com/download', + expiresAt: new Date(Date.now() + 60000).toISOString(), + fileSize: 1024, + fileName: 'metric-456.csv', + rowCount: 100, + }, + } as any); + + const result = await downloadMetricFileHandler(mockMetricId, mockUser as any); + + // Verify permission check was performed + expect(checkPermission).toHaveBeenCalledWith({ + userId: mockUser.id, + assetId: mockMetricId, + assetType: 'metric_file', + requiredRole: 'can_view', + organizationId: mockOrganizationId, + }); + + // Verify task was triggered with correct parameters + expect(tasks.trigger).toHaveBeenCalledWith('export-metric-data', { + metricId: mockMetricId, + userId: mockUser.id, + organizationId: mockOrganizationId, + }); + + // Verify successful response + expect(result).toMatchObject({ + downloadUrl: 'https://example.com/download', + fileSize: 1024, + fileName: 'metric-456.csv', + rowCount: 100, + message: 'Download link expires in 60 seconds. Please start your download immediately.', + }); + }); + + it('should respect different permission levels', async () => { + vi.mocked(getUserOrganizationId).mockResolvedValue({ + organizationId: mockOrganizationId, + userId: mockUser.id, + role: 'member', + } as any); + + // Test with 'owner' permission + vi.mocked(checkPermission).mockResolvedValue({ + hasAccess: true, + effectiveRole: 'owner', + accessPath: 'direct', + }); + + vi.mocked(tasks.trigger).mockResolvedValue({ id: 'task-123' } as any); + + const { runs } = await import('@trigger.dev/sdk'); + vi.mocked(runs.retrieve).mockResolvedValue({ + status: 'COMPLETED', + output: { + success: true, + downloadUrl: 'https://example.com/download', + expiresAt: new Date(Date.now() + 60000).toISOString(), + }, + } as any); + + await downloadMetricFileHandler(mockMetricId, mockUser as any); + + expect(checkPermission).toHaveBeenCalledWith({ + userId: mockUser.id, + assetId: mockMetricId, + assetType: 'metric_file', + requiredRole: 'can_view', + organizationId: mockOrganizationId, + }); + }); + + it('should handle workspace sharing permissions', async () => { + vi.mocked(getUserOrganizationId).mockResolvedValue({ + organizationId: mockOrganizationId, + userId: mockUser.id, + role: 'member', + } as any); + + // Simulate access through workspace sharing + vi.mocked(checkPermission).mockResolvedValue({ + hasAccess: true, + effectiveRole: 'can_view', + accessPath: 'workspace_sharing', + }); + + vi.mocked(tasks.trigger).mockResolvedValue({ id: 'task-123' } as any); + + const { runs } = await import('@trigger.dev/sdk'); + vi.mocked(runs.retrieve).mockResolvedValue({ + status: 'COMPLETED', + output: { + success: true, + downloadUrl: 'https://example.com/download', + expiresAt: new Date(Date.now() + 60000).toISOString(), + }, + } as any); + + await downloadMetricFileHandler(mockMetricId, mockUser as any); + + expect(checkPermission).toHaveBeenCalled(); + expect(tasks.trigger).toHaveBeenCalled(); + }); + + it('should handle cascading permissions', async () => { + vi.mocked(getUserOrganizationId).mockResolvedValue({ + organizationId: mockOrganizationId, + userId: mockUser.id, + role: 'member', + } as any); + + // Simulate access through cascading permissions (e.g., from parent collection) + vi.mocked(checkPermission).mockResolvedValue({ + hasAccess: true, + effectiveRole: 'can_edit', + accessPath: 'cascading', + }); + + vi.mocked(tasks.trigger).mockResolvedValue({ id: 'task-123' } as any); + + const { runs } = await import('@trigger.dev/sdk'); + vi.mocked(runs.retrieve).mockResolvedValue({ + status: 'COMPLETED', + output: { + success: true, + downloadUrl: 'https://example.com/download', + expiresAt: new Date(Date.now() + 60000).toISOString(), + }, + } as any); + + await downloadMetricFileHandler(mockMetricId, mockUser as any); + + expect(checkPermission).toHaveBeenCalled(); + expect(tasks.trigger).toHaveBeenCalled(); + }); + }); + + describe('Error Handling', () => { + beforeEach(() => { + // Setup basic mocks for permission checks to pass + vi.mocked(getUserOrganizationId).mockResolvedValue({ + organizationId: mockOrganizationId, + userId: mockUser.id, + role: 'member', + } as any); + + vi.mocked(checkPermission).mockResolvedValue({ + hasAccess: true, + effectiveRole: 'can_view', + accessPath: 'direct', + }); + }); + + it('should handle task failure', async () => { + vi.mocked(tasks.trigger).mockResolvedValue({ id: 'task-123' } as any); + + const { runs } = await import('@trigger.dev/sdk'); + vi.mocked(runs.retrieve).mockResolvedValue({ + status: 'FAILED', + } as any); + + await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toThrow( + HTTPException + ); + + await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toMatchObject({ + status: 500, + message: 'Export task failed', + }); + }); + + it('should handle timeout', async () => { + vi.mocked(tasks.trigger).mockResolvedValue({ id: 'task-123' } as any); + + const { runs } = await import('@trigger.dev/sdk'); + + // Simulate task still in progress after timeout + let callCount = 0; + vi.mocked(runs.retrieve).mockImplementation(async () => { + callCount++; + // Return EXECUTING status to simulate timeout + return { status: 'EXECUTING' } as any; + }); + + // Speed up test by reducing timeout + vi.useFakeTimers(); + + const promise = downloadMetricFileHandler(mockMetricId, mockUser as any); + + // Fast-forward time to trigger timeout + await vi.runAllTimersAsync(); + + await expect(promise).rejects.toThrow(HTTPException); + + vi.useRealTimers(); + }); + + it('should handle specific error codes from export task', async () => { + vi.mocked(tasks.trigger).mockResolvedValue({ id: 'task-123' } as any); + + const { runs } = await import('@trigger.dev/sdk'); + + // Test UNAUTHORIZED error code + vi.mocked(runs.retrieve).mockResolvedValue({ + status: 'COMPLETED', + output: { + success: false, + error: 'Access denied to data source', + errorCode: 'UNAUTHORIZED', + }, + } as any); + + await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toMatchObject({ + status: 403, + message: 'Access denied to data source', + }); + + // Test NOT_FOUND error code + vi.mocked(runs.retrieve).mockResolvedValue({ + status: 'COMPLETED', + output: { + success: false, + error: 'Metric not found', + errorCode: 'NOT_FOUND', + }, + } as any); + + await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toMatchObject({ + status: 404, + message: 'Metric file not found or data source credentials missing', + }); + + // Test QUERY_ERROR + vi.mocked(runs.retrieve).mockResolvedValue({ + status: 'COMPLETED', + output: { + success: false, + error: 'SQL syntax error', + errorCode: 'QUERY_ERROR', + }, + } as any); + + await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toMatchObject({ + status: 400, + message: 'Query execution failed: SQL syntax error', + }); + }); + }); +}); diff --git a/apps/server/src/api/v2/metric_files/download-metric-file.ts b/apps/server/src/api/v2/metric_files/download-metric-file.ts index 183c54705..816ace4fd 100644 --- a/apps/server/src/api/v2/metric_files/download-metric-file.ts +++ b/apps/server/src/api/v2/metric_files/download-metric-file.ts @@ -1,3 +1,4 @@ +import { type AssetPermissionCheck, checkPermission } from '@buster/access-controls'; import type { User } from '@buster/database'; import { getUserOrganizationId } from '@buster/database'; import type { ExportMetricDataOutput, MetricDownloadResponse } from '@buster/server-shared/metrics'; @@ -9,9 +10,10 @@ import { HTTPException } from 'hono/http-exception'; * * This handler: * 1. Validates user has access to the organization - * 2. Triggers the export task in Trigger.dev - * 3. Waits for the task to complete (max 2 minutes) - * 4. Returns a presigned URL for downloading the CSV file + * 2. Checks user has permission to view the metric file + * 3. Triggers the export task in Trigger.dev + * 4. Waits for the task to complete (max 2 minutes) + * 5. Returns a presigned URL for downloading the CSV file * * The download URL expires after 60 seconds for security */ @@ -30,6 +32,23 @@ export async function downloadMetricFileHandler( const { organizationId } = userOrg; + // Check if user has permission to view this metric file + const permissionCheck: AssetPermissionCheck = { + userId: user.id, + assetId: metricId, + assetType: 'metric_file', + requiredRole: 'can_view', + organizationId, + }; + + const permissionResult = await checkPermission(permissionCheck); + + if (!permissionResult.hasAccess) { + throw new HTTPException(403, { + message: 'You do not have permission to download this metric file', + }); + } + try { // Trigger the export task const handle = await tasks.trigger('export-metric-data', { diff --git a/apps/trigger/src/tasks/export-metric-data/export-metric-data.ts b/apps/trigger/src/tasks/export-metric-data/export-metric-data.ts index d5987eb92..ddf92dcb3 100644 --- a/apps/trigger/src/tasks/export-metric-data/export-metric-data.ts +++ b/apps/trigger/src/tasks/export-metric-data/export-metric-data.ts @@ -1,6 +1,7 @@ import { randomBytes } from 'node:crypto'; import { GetObjectCommand, PutObjectCommand, S3Client } from '@aws-sdk/client-s3'; import { getSignedUrl } from '@aws-sdk/s3-request-presigner'; +import { type AssetPermissionCheck, checkPermission } from '@buster/access-controls'; import { createAdapter } from '@buster/data-source'; import type { Credentials } from '@buster/data-source'; import { getDataSourceCredentials, getMetricForExport } from '@buster/database'; @@ -41,13 +42,14 @@ const MAX_FILE_SIZE = 500 * 1024 * 1024; // 500MB max file size * Task for exporting metric data to CSV and generating a presigned download URL * * This task: - * 1. Fetches metric configuration and validates user access - * 2. Retrieves data source credentials from vault - * 3. Executes the metric's SQL query - * 4. Converts results to CSV format - * 5. Uploads to R2 storage - * 6. Generates a 60-second presigned URL for download - * 7. Schedules cleanup after 60 seconds + * 1. Fetches metric configuration and validates organization access + * 2. Checks user has permission to view the metric file + * 3. Retrieves data source credentials from vault + * 4. Executes the metric's SQL query + * 5. Converts results to CSV format + * 6. Uploads to R2 storage + * 7. Generates a 60-second presigned URL for download + * 8. Schedules cleanup after 60 seconds */ export const exportMetricData: ReturnType< typeof schemaTask< @@ -96,6 +98,31 @@ export const exportMetricData: ReturnType< }; } + // Step 1b: Check user has permission to view this metric file + const permissionCheck: AssetPermissionCheck = { + userId: payload.userId, + assetId: payload.metricId, + assetType: 'metric_file', + requiredRole: 'can_view', + organizationId: payload.organizationId, + }; + + const permissionResult = await checkPermission(permissionCheck); + + if (!permissionResult.hasAccess) { + logger.error('User lacks permission to access metric file', { + metricId: payload.metricId, + userId: payload.userId, + organizationId: payload.organizationId, + }); + + return { + success: false, + error: 'You do not have permission to export this metric', + errorCode: 'UNAUTHORIZED', + }; + } + logger.log('Metric validated', { metricName: metric.name, dataSourceId: metric.dataSourceId, diff --git a/packages/access-controls/src/types/asset-permissions.ts b/packages/access-controls/src/types/asset-permissions.ts index 333005601..7f8769881 100644 --- a/packages/access-controls/src/types/asset-permissions.ts +++ b/packages/access-controls/src/types/asset-permissions.ts @@ -14,6 +14,7 @@ export const AssetTypeSchema = z.enum([ 'chat', 'metric_file', 'dashboard_file', + 'report_file', 'data_source', 'metric', 'filter', @@ -32,6 +33,7 @@ export type CascadingAssetType = | 'chat' | 'metric_file' | 'dashboard_file' + | 'report_file' | 'collection'; // Permission roles - matching database AssetPermissionRole enum diff --git a/packages/ai/src/agents/think-and-prep-agent/investigation-instructions.ts b/packages/ai/src/agents/think-and-prep-agent/investigation-instructions.ts index feaedf1be..6ca9bab67 100644 --- a/packages/ai/src/agents/think-and-prep-agent/investigation-instructions.ts +++ b/packages/ai/src/agents/think-and-prep-agent/investigation-instructions.ts @@ -715,11 +715,11 @@ export const getThinkAndPrepInstructions = async ({ const userId = runtimeContext.get('userId'); const dataSourceSyntax = runtimeContext.get('dataSourceSyntax'); - const datasets = await getPermissionedDatasets(userId, 0, 1000); + const datasetResults = await getPermissionedDatasets({ userId, page: 0, pageSize: 1000 }); // Extract yml_content from each dataset and join with separators - const assembledYmlContent = datasets - .map((dataset: { ymlFile: string | null | undefined }) => dataset.ymlFile) + const assembledYmlContent = datasetResults.datasets + .map((dataset: { ymlContent: string | null | undefined }) => dataset.ymlContent) .filter((content: string | null | undefined) => content !== null && content !== undefined) .join('\n---\n'); diff --git a/packages/database/src/queries/asset-permissions/check-asset-permission.ts b/packages/database/src/queries/asset-permissions/check-asset-permission.ts index 42c22f397..9be02a11c 100644 --- a/packages/database/src/queries/asset-permissions/check-asset-permission.ts +++ b/packages/database/src/queries/asset-permissions/check-asset-permission.ts @@ -17,7 +17,14 @@ import type { AssetPermissionRole, AssetType, WorkspaceSharing } from '../../sch export interface CheckAssetPermissionParams { userId: string; assetId: string; - assetType: 'dashboard' | 'thread' | 'chat' | 'metric_file' | 'dashboard_file' | 'collection'; + assetType: + | 'dashboard' + | 'thread' + | 'chat' + | 'metric_file' + | 'dashboard_file' + | 'report_file' + | 'collection'; organizationId?: string; } diff --git a/packages/database/src/schema-types/enums.ts b/packages/database/src/schema-types/enums.ts index d92d6a2a9..599b82e1c 100644 --- a/packages/database/src/schema-types/enums.ts +++ b/packages/database/src/schema-types/enums.ts @@ -6,6 +6,7 @@ export type AssetType = | 'chat' | 'metric_file' | 'dashboard_file' + | 'report_file' | 'collection' | 'data_source' | 'metric' diff --git a/packages/server-shared/src/index.ts b/packages/server-shared/src/index.ts index 1a1b9c48d..7fedc338a 100644 --- a/packages/server-shared/src/index.ts +++ b/packages/server-shared/src/index.ts @@ -12,6 +12,8 @@ export * from './message'; export * from './metrics'; export * from './organization'; export * from './security'; +// Export share module (has some naming conflicts with chats and metrics) +// TODO: Resolve naming conflicts properly export * from './share'; export * from './slack'; export * from './teams'; diff --git a/packages/server-shared/src/share/assets.ts b/packages/server-shared/src/share/assets.ts index 3849b7552..ef07581c2 100644 --- a/packages/server-shared/src/share/assets.ts +++ b/packages/server-shared/src/share/assets.ts @@ -12,6 +12,9 @@ const AssetPermissionRoleEnums: Record; diff --git a/packages/server-shared/src/share/requests.ts b/packages/server-shared/src/share/requests.ts index ac3e07b29..2ecff914d 100644 --- a/packages/server-shared/src/share/requests.ts +++ b/packages/server-shared/src/share/requests.ts @@ -13,7 +13,7 @@ export const SharePostRequestSchema = z.array( export type SharePostRequest = z.infer; //Used for updating share permissions for a report, collection, or metric -export const ShareUpdateRequestSchema = z.object({ +export const SharePermissionsUpdateRequestSchema = z.object({ publicly_accessible: z.boolean().optional(), public_expiry_date: z.string().nullable().optional(), public_password: z.string().nullable().optional(), @@ -28,7 +28,11 @@ export const ShareUpdateRequestSchema = z.object({ .optional(), }); -export type ShareUpdateRequest = z.infer; +export type SharePermissionsUpdateRequest = z.infer; + +// Keep old names for backward compatibility but don't export from index +const ShareUpdateRequestSchema = SharePermissionsUpdateRequestSchema; +type ShareUpdateRequest = SharePermissionsUpdateRequest; //Used for deleting share permissions for a report, collection, or metric export const ShareDeleteRequestSchema = z.array(z.string());