Merge pull request #463 from levi0005/fix/key-management-improvements

Fix/key management improvements

Author: Jambox11

src/key-management/events/key-generated.event.ts

@@ -0,0 +1,7 @@
+export class KeyGeneratedEvent {
+  constructor(
+    public readonly publicKey: string,
+    public readonly keyType: string,
+    public readonly timestamp: Date,
+  ) {}
+}

src/key-management/events/key-rotated.event.ts

@@ -0,0 +1,8 @@
+export class KeyRotatedEvent {
+  constructor(
+    public readonly predecessorWalletId: string,
+    public readonly successorWalletId: string,
+    public readonly successorPublicKey: string,
+    public readonly timestamp: Date,
+  ) {}
+}

src/key-management/events/key-signed.event.ts

@@ -0,0 +1,6 @@
+export class KeySignedEvent {
+  constructor(
+    public readonly publicKey: string,
+    public readonly timestamp: Date,
+  ) {}
+}

src/key-management/events/key-validated.event.ts

@@ -0,0 +1,8 @@
+export class KeyValidatedEvent {
+  constructor(
+    public readonly publicKey: string,
+    public readonly keyType: string,
+    public readonly valid: boolean,
+    public readonly timestamp: Date,
+  ) {}
+}

src/key-management/key-management.controller.spec.ts

@@ -1,4 +1,5 @@
 import { Test, TestingModule } from '@nestjs/testing';
+import { BadRequestException } from '@nestjs/common';
 import { KeyManagementController } from './key-management.controller';
 import { KeyManagementService } from './key-management.service';
 import { KeyType } from './domain/key-types';
@@ -113,33 +114,100 @@ describe('KeyManagementController', () => {
     });
   });
 
+  describe('generateKey input validation', () => {
+    it('should throw BadRequestException when keyType is missing', async () => {
+      await expect(
+        controller.generateKey({ keyType: undefined as any }),
+      ).rejects.toThrow(BadRequestException);
+    });
+
+    it('should throw BadRequestException when keyType is invalid', async () => {
+      await expect(
+        controller.generateKey({ keyType: 'BOGUS' as any }),
+      ).rejects.toThrow(BadRequestException);
+    });
+  });
+
   describe('getAuditLog', () => {
-    it('should return audit logs with default limit', async () => {
-      const mockLogs = [
+    const mockResult = {
+      data: [
         {
           operation: 'GENERATE',
           keyId: 'key-1',
           publicKey: 'GPUBLIC123...',
           timestamp: new Date(),
           success: true,
         },
-      ];
+      ],
+      total: 1,
+      limit: 100,
+      offset: 0,
+      hasMore: false,
+    };
 
-      mockKeyManagementService.getAuditLog.mockReturnValue(mockLogs);
+    it('should return paginated audit logs with default params', async () => {
+      mockKeyManagementService.getAuditLog.mockReturnValue(mockResult);
 
       const result = await controller.getAuditLog();
 
-      expect(result).toEqual({ logs: mockLogs });
-      expect(service.getAuditLog).toHaveBeenCalledWith(100);
+      expect(result).toEqual({
+        logs: mockResult.data,
+        total: mockResult.total,
+        limit: mockResult.limit,
+        offset: mockResult.offset,
+        hasMore: mockResult.hasMore,
+      });
+      expect(service.getAuditLog).toHaveBeenCalledWith(
+        expect.objectContaining({ limit: undefined, offset: undefined }),
+      );
+    });
+
+    it('should pass limit and offset to service', async () => {
+      mockKeyManagementService.getAuditLog.mockReturnValue(mockResult);
+
+      await controller.getAuditLog(undefined, undefined, undefined, undefined, undefined, '50', '10');
+
+      expect(service.getAuditLog).toHaveBeenCalledWith(
+        expect.objectContaining({ limit: 50, offset: 10 }),
+      );
     });
 
-    it('should return audit logs with custom limit', async () => {
-      const mockLogs = [];
-      mockKeyManagementService.getAuditLog.mockReturnValue(mockLogs);
+    it('should pass operation filter to service', async () => {
+      mockKeyManagementService.getAuditLog.mockReturnValue(mockResult);
+
+      await controller.getAuditLog('GENERATE');
+
+      expect(service.getAuditLog).toHaveBeenCalledWith(
+        expect.objectContaining({ operation: 'GENERATE' }),
+      );
+    });
+
+    it('should parse success filter', async () => {
+      mockKeyManagementService.getAuditLog.mockReturnValue(mockResult);
+
+      await controller.getAuditLog(undefined, undefined, 'true');
 
-      await controller.getAuditLog('50');
+      expect(service.getAuditLog).toHaveBeenCalledWith(
+        expect.objectContaining({ success: true }),
+      );
+    });
+
+    it('should throw BadRequestException for non-integer limit', async () => {
+      await expect(
+        controller.getAuditLog(undefined, undefined, undefined, undefined, undefined, 'abc'),
+      ).rejects.toThrow(BadRequestException);
+    });
+
+    it('should throw BadRequestException for limit exceeding 100', async () => {
+      await expect(
+        controller.getAuditLog(undefined, undefined, undefined, undefined, undefined, '200'),
+      ).rejects.toThrow(BadRequestException);
+    });
 
-      expect(service.getAuditLog).toHaveBeenCalledWith(50);
+    it('should throw BadRequestException for invalid startDate', async () => {
+      await expect(
+        controller.getAuditLog(undefined, undefined, undefined, 'not-a-date'),
+      ).rejects.toThrow(BadRequestException);
     });
   });
 

src/key-management/key-management.controller.ts

@@ -1,4 +1,5 @@
 import {
+  BadRequestException,
   Controller,
   Post,
   Body,
@@ -26,6 +27,31 @@ import {
 } from './key-rotation-audit.service';
 import { KeyOperation } from '../generated/prisma/client';
 
+function parsePaginationParam(
+  value: string | undefined,
+  name: string,
+  max = 100,
+): number | undefined {
+  if (value === undefined) return undefined;
+  const n = Number(value);
+  if (!Number.isInteger(n) || n < 0) {
+    throw new BadRequestException(`${name} must be a non-negative integer`);
+  }
+  if (name === 'limit' && n > max) {
+    throw new BadRequestException(`limit must not exceed ${max}`);
+  }
+  return n;
+}
+
+function parseDate(value: string | undefined, name: string): Date | undefined {
+  if (value === undefined) return undefined;
+  const d = new Date(value);
+  if (isNaN(d.getTime())) {
+    throw new BadRequestException(`${name} must be a valid ISO date string`);
+  }
+  return d;
+}
+
 /**
  * Internal controller for key management operations
  *
@@ -50,6 +76,15 @@ export class KeyManagementController {
   @Post('generate')
   @HttpCode(HttpStatus.OK)
   async generateKey(@Body() request: GenerateKeyRequest) {
+    if (!request?.keyType) {
+      throw new BadRequestException('keyType is required');
+    }
+    if (!Object.values(KeyType).includes(request.keyType)) {
+      throw new BadRequestException(
+        `Invalid keyType: "${request.keyType}". Must be one of: ${Object.values(KeyType).join(', ')}`,
+      );
+    }
+
     const result = await this.keyManagementService.generateKey(request);
 
     return {
@@ -133,17 +168,49 @@ export class KeyManagementController {
   }
 
   /**
-   * Gets audit log (admin only)
+   * Gets in-memory audit log with optional filtering and pagination
+   *
+   * Query parameters:
+   * - operation: Filter by operation type (GENERATE, SIGN, ROTATE, etc.)
+   * - publicKey: Filter by public key
+   * - success: Filter by success status (true/false)
+   * - startDate: Start of date range (ISO string)
+   * - endDate: End of date range (ISO string)
+   * - limit: Max results (default: 100, max: 100)
+   * - offset: Pagination offset (default: 0)
    */
   @ApiOperation({ summary: 'Retrieve in-memory audit log (internal, admin only)' })
   @ApiQuery({ name: 'limit', required: false, description: 'Max entries to return (default: 100)' })
   @ApiResponse({ status: 200, description: 'Array of audit log entries' })
   @Get('audit')
-  async getAuditLog(@Query('limit') limit?: string) {
-    const auditLimit = limit ? parseInt(limit, 10) : 100;
-    const logs = this.keyManagementService.getAuditLog(auditLimit);
+  async getAuditLog(
+    @Query('operation') operation?: string,
+    @Query('publicKey') publicKey?: string,
+    @Query('success') success?: string,
+    @Query('startDate') startDate?: string,
+    @Query('endDate') endDate?: string,
+    @Query('limit') limit?: string,
+    @Query('offset') offset?: string,
+  ) {
+    const query: AuditLogQuery = {
+      operation,
+      publicKey,
+      success: success !== undefined ? success === 'true' : undefined,
+      startDate: parseDate(startDate, 'startDate'),
+      endDate: parseDate(endDate, 'endDate'),
+      limit: parsePaginationParam(limit, 'limit'),
+      offset: parsePaginationParam(offset, 'offset'),
+    };
+
+    const result = this.keyManagementService.getAuditLog(query);
 
-    return { logs };
+    return {
+      logs: result.data,
+      total: result.total,
+      limit: result.limit,
+      offset: result.offset,
+      hasMore: result.hasMore,
+    };
   }
 
   /**
@@ -168,8 +235,8 @@ export class KeyManagementController {
     @Query('operation') operation?: string,
   ) {
     const query: KeyStatisticsQuery = {
-      startDate: startDate ? new Date(startDate) : undefined,
-      endDate: endDate ? new Date(endDate) : undefined,
+      startDate: parseDate(startDate, 'startDate'),
+      endDate: parseDate(endDate, 'endDate'),
       operation,
     };
 
@@ -203,8 +270,8 @@ export class KeyManagementController {
     @Query('includeTimeSeries') includeTimeSeries?: string,
   ) {
     const query: KeyStatisticsQuery = {
-      startDate: startDate ? new Date(startDate) : undefined,
-      endDate: endDate ? new Date(endDate) : undefined,
+      startDate: parseDate(startDate, 'startDate'),
+      endDate: parseDate(endDate, 'endDate'),
       operation,
       includeTimeSeries: includeTimeSeries === 'true',
     };

src/key-management/key-management.module.ts

@@ -9,7 +9,7 @@ import { PrismaModule } from '../prisma/prisma.module';
 import { KeyManagementMetricsService } from './key-management-metrics.service';
 
 @Module({
-  imports: [EncryptionModule, PrismaModule],
+  imports: [EncryptionModule, PrismaModule, EventEmitterModule.forRoot()],
   controllers: [KeyManagementController],
   providers: [
     KeyManagementService,