feat(protect): init audit logging

Author: calvinbrewer

packages/protect/__tests__/protect.test.ts

@@ -1,5 +1,5 @@
 import 'dotenv/config'
-import { describe, expect, it, vi } from 'vitest'
+import { describe, expect, it, beforeAll } from 'vitest'
 
 import { LockContext, protect, csTable, csColumn } from '../src'
 
@@ -17,10 +17,16 @@ type User = {
   number?: number
 }
 
+let protectClient: Awaited<ReturnType<typeof protect>>
+
+beforeAll(async () => {
+  protectClient = await protect({
+    schemas: [users],
+  })
+})
+
 describe('encryption and decryption', () => {
   it('should encrypt and decrypt a payload', async () => {
-    const protectClient = await protect({ schemas: [users] })
-
     const email = 'hello@example.com'
 
     const ciphertext = await protectClient.encrypt(email, {
@@ -35,16 +41,18 @@ describe('encryption and decryption', () => {
     // Verify encrypted field
     expect(ciphertext.data).toHaveProperty('c')
 
-    const plaintext = await protectClient.decrypt(ciphertext.data)
+    const plaintext = await protectClient.decrypt(ciphertext.data).audit({
+      metadata: {
+        user: 'cj@cjb.io',
+      },
+    })
 
     expect(plaintext).toEqual({
       data: email,
     })
   }, 30000)
 
   it('should return null if plaintext is null', async () => {
-    const protectClient = await protect({ schemas: [users] })
-
     const ciphertext = await protectClient.encrypt(null, {
       column: users.email,
       table: users,
@@ -65,8 +73,6 @@ describe('encryption and decryption', () => {
   }, 30000)
 
   it('should encrypt and decrypt a model', async () => {
-    const protectClient = await protect({ schemas: [users] })
-
     // Create a model with decrypted values
     const decryptedModel = {
       id: '1',
@@ -117,8 +123,6 @@ describe('encryption and decryption', () => {
   }, 30000)
 
   it('should handle null values in a model', async () => {
-    const protectClient = await protect({ schemas: [users] })
-
     // Create a model with null values
     const decryptedModel = {
       id: '1',
@@ -169,8 +173,6 @@ describe('encryption and decryption', () => {
   }, 30000)
 
   it('should handle undefined values in a model', async () => {
-    const protectClient = await protect({ schemas: [users] })
-
     // Create a model with undefined values
     const decryptedModel = {
       id: '1',
@@ -223,8 +225,6 @@ describe('encryption and decryption', () => {
 
 describe('bulk encryption', () => {
   it('should bulk encrypt and decrypt models', async () => {
-    const protectClient = await protect({ schemas: [users] })
-
     // Create models with decrypted values
     const decryptedModels = [
       {
@@ -301,8 +301,6 @@ describe('bulk encryption', () => {
   }, 30000)
 
   it('should return empty array if models is empty', async () => {
-    const protectClient = await protect({ schemas: [users] })
-
     // Encrypt empty array of models
     const encryptedModels = await protectClient.bulkEncryptModels<User>(
       [],
@@ -317,8 +315,6 @@ describe('bulk encryption', () => {
   }, 30000)
 
   it('should return empty array if decrypting empty array of models', async () => {
-    const protectClient = await protect({ schemas: [users] })
-
     // Decrypt empty array of models
     const decryptedResult = await protectClient.bulkDecryptModels<User>([])
 
@@ -332,7 +328,6 @@ describe('bulk encryption', () => {
 
 describe('bulk encryption edge cases', () => {
   it('should handle mixed null and non-null values in bulk operations', async () => {
-    const protectClient = await protect({ schemas: [users] })
     const decryptedModels = [
       {
         id: '1',
@@ -405,7 +400,6 @@ describe('bulk encryption edge cases', () => {
   }, 30000)
 
   it('should handle mixed undefined and non-undefined values in bulk operations', async () => {
-    const protectClient = await protect({ schemas: [users] })
     const decryptedModels = [
       {
         id: '1',
@@ -478,7 +472,6 @@ describe('bulk encryption edge cases', () => {
   }, 30000)
 
   it('should handle empty models in bulk operations', async () => {
-    const protectClient = await protect({ schemas: [users] })
     const decryptedModels = [
       {
         id: '1',
@@ -548,7 +541,6 @@ describe('bulk encryption edge cases', () => {
 
 describe('error handling', () => {
   it('should handle invalid encrypted payloads', async () => {
-    const protectClient = await protect({ schemas: [users] })
     const validModel = {
       id: '1',
       email: 'test@example.com',
@@ -582,7 +574,6 @@ describe('error handling', () => {
   }, 30000)
 
   it('should handle missing required fields', async () => {
-    const protectClient = await protect({ schemas: [users] })
     const model = {
       id: '1',
       email: null,
@@ -603,7 +594,6 @@ describe('error handling', () => {
 
 describe('type safety', () => {
   it('should maintain type safety with complex nested objects', async () => {
-    const protectClient = await protect({ schemas: [users] })
     const model = {
       id: '1',
       email: 'test@example.com',
@@ -641,7 +631,6 @@ describe('type safety', () => {
 
 describe('performance', () => {
   it('should handle large numbers of models efficiently', async () => {
-    const protectClient = await protect({ schemas: [users] })
     const largeModels = Array(10)
       .fill(null)
       .map((_, i) => ({

packages/protect/src/ffi/operations/base-operation.ts

@@ -0,0 +1,24 @@
+export type AuditConfig = {
+  metadata?: Record<string, unknown>
+}
+
+export class BaseOperation<T> {
+  protected auditMetadata?: Record<string, unknown>
+
+  /**
+   * Attach audit metadata to this operation. Can be chained.
+   * @param config Configuration for ZeroKMS audit logging
+   * @param config.metadata Arbitrary JSON object for appending metadata to the audit log
+   */
+  audit(config: AuditConfig): this {
+    this.auditMetadata = config.metadata
+    return this
+  }
+
+  /**
+   * Get the audit metadata for this operation.
+   */
+  public getAuditMetadata(): Record<string, unknown> | undefined {
+    return this.auditMetadata
+  }
+}

packages/protect/src/ffi/operations/decrypt.ts

@@ -5,22 +5,29 @@ import { type ProtectError, ProtectErrorTypes } from '../..'
 import { logger } from '../../../../utils/logger'
 import type { LockContext } from '../../identify'
 import type { Client, EncryptedPayload } from '../../types'
+import { BaseOperation } from './base-operation'
 
 export class DecryptOperation
+  extends BaseOperation<Result<string | null, ProtectError>>
   implements PromiseLike<Result<string | null, ProtectError>>
 {
   private client: Client
   private encryptedData: EncryptedPayload
 
   constructor(client: Client, encryptedData: EncryptedPayload) {
+    super()
     this.client = client
     this.encryptedData = encryptedData
   }
 
   public withLockContext(
     lockContext: LockContext,
   ): DecryptOperationWithLockContext {
-    return new DecryptOperationWithLockContext(this, lockContext)
+    const opWithLock = new DecryptOperationWithLockContext(this, lockContext)
+    if (this.getAuditMetadata()) {
+      opWithLock.audit(this.getAuditMetadata() ?? {})
+    }
+    return opWithLock
   }
 
   public then<TResult1 = Result<string | null, ProtectError>, TResult2 = never>(
@@ -46,7 +53,9 @@ export class DecryptOperation
           return null
         }
 
-        logger.debug('Decrypting data WITHOUT a lock context')
+        logger.debug('Decrypting data WITHOUT a lock context', {
+          auditMetadata: this.getAuditMetadata(),
+        })
         return await ffiDecrypt(this.client, this.encryptedData.c)
       },
       (error) => ({
@@ -59,23 +68,30 @@ export class DecryptOperation
   public getOperation(): {
     client: Client
     encryptedData: EncryptedPayload
+    auditMetadata?: Record<string, unknown>
   } {
     return {
       client: this.client,
       encryptedData: this.encryptedData,
+      auditMetadata: this.getAuditMetadata(),
     }
   }
 }
 
 export class DecryptOperationWithLockContext
+  extends BaseOperation<Result<string | null, ProtectError>>
   implements PromiseLike<Result<string | null, ProtectError>>
 {
   private operation: DecryptOperation
   private lockContext: LockContext
 
   constructor(operation: DecryptOperation, lockContext: LockContext) {
+    super()
     this.operation = operation
     this.lockContext = lockContext
+    if (operation.getAuditMetadata()) {
+      this.audit(operation.getAuditMetadata() ?? {})
+    }
   }
 
   public then<TResult1 = Result<string | null, ProtectError>, TResult2 = never>(
@@ -103,7 +119,9 @@ export class DecryptOperationWithLockContext
           return null
         }
 
-        logger.debug('Decrypting data WITH a lock context')
+        logger.debug('Decrypting data WITH a lock context', {
+          auditMetadata: this.getAuditMetadata(),
+        })
 
         const context = await this.lockContext.getLockContext()