Adding KMS keyring

apps/policy-engine/.env.default

@@ -7,3 +7,7 @@ POLICY_ENGINE_DATABASE_URL="postgresql://postgres:postgres@localhost:5432/engine
 ENGINE_UID="local-dev-engine-instance-1"
 
 MASTER_PASSWORD="unsafe-local-dev-master-password"
+
+KEYRING_TYPE="raw"
+
+# MASTER_AWS_KMS_ARN="arn:aws:kms:us-east-2:728783560968:key/f6aa3ddb-47c3-4f31-977d-b93205bb23d1"

apps/policy-engine/.env.test.default

@@ -1,7 +1,13 @@
-# IMPORTANT: The variables defined here will override other variables.
-# See `./apps/policy-engine/jest.setup.ts`.
 NODE_ENV=test
 
+PORT=3010
+
 POLICY_ENGINE_DATABASE_URL="postgresql://postgres:postgres@localhost:5432/engine-test?schema=public"
 
-MASTER_PASSWORD="unsafe-local-dev-master-password"
+ENGINE_UID="local-dev-engine-instance-1"
+
+MASTER_PASSWORD="unsafe-local-test-master-password"
+
+KEYRING_TYPE="raw"
+
+# MASTER_AWS_KMS_ARN="arn:aws:kms:us-east-2:728783560968:key/f6aa3ddb-47c3-4f31-977d-b93205bb23d1"

apps/policy-engine/jest.setup.ts

@@ -3,7 +3,6 @@ import fs from 'fs'
 import nock from 'nock'
 
 const testEnvFile = `${__dirname}/.env.test`
-const envFile = `${__dirname}/.env`
 
 // Ensure a test environment variable file exists because of the override config
 // loading mechanics below.
@@ -17,7 +16,15 @@ if (!fs.existsSync(testEnvFile)) {
 // If a .env.test file is not found, the DATABASE_URL will fallback to the
 // default. Consequently, you'll lose your development database during the
 // integration tests teardown. Hence, the check above.
-dotenv.config({ path: envFile })
+
+// We also don't even want any .env values; just kill them during tests.
+// Clear process.env
+for (const prop in process.env) {
+  if (Object.prototype.hasOwnProperty.call(process.env, prop)) {
+    delete process.env[prop]
+  }
+}
+
 dotenv.config({ path: testEnvFile, override: true })
 
 // Disable outgoing HTTP requests to avoid flaky tests.

apps/policy-engine/src/encryption/core/__test__/unit/encryption.service.spec.ts

@@ -1,29 +1,39 @@
+import { ConfigModule, ConfigService } from '@nestjs/config'
 import { Test } from '@nestjs/testing'
-
-import { ConfigModule } from '@nestjs/config'
+import { mock } from 'jest-mock-extended'
+import nock from 'nock'
 import { load } from '../../../../policy-engine.config'
 import { EncryptionRepository } from '../../../persistence/repository/encryption.repository'
 import { EncryptionService } from '../../encryption.service'
 
 describe('EncryptionService', () => {
   let service: EncryptionService
 
+  nock.enableNetConnect('kms.us-east-2.amazonaws.com:443')
+
   beforeEach(async () => {
-    // Create a mock ConfigService
-    const configModuleMock = {
+    // These mocked config values matter; they're specifically tied to the mocked masterKey below
+    // If you change these, the decryption won't work & tests will fail
+    const configServiceMock = mock<ConfigService>({
       get: jest.fn().mockImplementation((key: string) => {
-        if (key === 'engine.masterPassword') {
-          return 'unsafe-local-dev-master-password'
+        if (key === 'keyring') {
+          return {
+            type: 'raw',
+            masterPassword: 'unsafe-local-dev-master-password'
+          }
+        }
+        if (key === 'engine.id') {
+          return 'local-dev-engine-instance-1'
         }
       })
-    }
+    })
+
     const encryptionRepositoryMock = {
       getEngine: jest.fn().mockImplementation(() =>
         Promise.resolve({
           // unencryptedMasterKey: dfd9cc70f1ad02d19e0efa020d82f557022f59ca6bedbec1df38e8fd37ae3bb9
           masterKey:
             '0205785d67737fa3bae8eb249cf8d3baed5942f1677d8c98b4cdeef55560a3bcf510bd008d00030003617070000d61726d6f72792d656e67696e6500156177732d63727970746f2d7075626c69632d6b657900444177336764324b6e58646f512f2b76745347367031444442384d65766d61434b324c7861426e65476a315531537777526b376b4d366868752f707a446f48724c77773d3d0007707572706f7365000f646174612d656e6372797074696f6e000100146e617276616c2e61726d6f72792e656e67696e65002561726d6f72792e656e67696e652e6b656b000000800000000c8a92a7c9deb43316f6c29e8d0030132d63c7337c9888a06b638966e83056a0575958b42588b7aed999b9659e6d4bc5bed4664d91fae0b14d48917e00cdbb02000010000749ed0ed3616b7990f9e73f5a42eb46dc182002612e33dcb8e3c7d4759184c46ce3f0893a87ac15257d53097ac5d74affffffff00000001000000000000000000000001000000205d7209b51db8cf8264b9065add71a8514dc26baa6987d8a0a3acb1c4a2503b0f3b7c974a35ed234c1b94668736cd8bfa00673065023100a5d8d192e9802649dab86af6e00ab6d7472533e85dfe1006cb8bd9ef2472d15096fa42e742d18cb92530c762c3bd44d40230350299b42feaa1149c6ad1b25add24c30b3bf1c08263b96df0d43e2ad3e19802872e792040f1faf3d0a73bca6fb067ca',
-          adminApiKey: 'hex-admin-api-key',
           id: 'test-engine-id'
         })
       )
@@ -42,8 +52,8 @@ describe('EncryptionService', () => {
           useValue: encryptionRepositoryMock
         },
         {
-          provide: ConfigModule,
-          useValue: configModuleMock // use the mock ConfigService
+          provide: ConfigService,
+          useValue: configServiceMock // use the mock ConfigService
         }
       ]
     }).compile()
@@ -69,6 +79,4 @@ describe('EncryptionService', () => {
 
     expect(decrypted.toString('hex')).toBe(data)
   })
-
-  // Add your tests here
 })

apps/policy-engine/src/encryption/core/encryption.service.ts

@@ -1,5 +1,6 @@
 import {
   CommitmentPolicy,
+  KmsKeyringNode,
   RawAesKeyringNode,
   RawAesWrappingSuiteIdentifier,
   buildClient
@@ -30,7 +31,7 @@ export class EncryptionService implements OnApplicationBootstrap {
 
   private masterKey: Buffer | undefined
 
-  private adminApiKey: Buffer | undefined
+  private keyring: RawAesKeyringNode | KmsKeyringNode | undefined
 
   constructor(
     private encryptionRepository: EncryptionRepository,
@@ -42,22 +43,43 @@ export class EncryptionService implements OnApplicationBootstrap {
 
   async onApplicationBootstrap(): Promise<void> {
     this.logger.log('Keyring Service boot')
-    let engine = await this.encryptionRepository.getEngine(this.engineId)
-
-    // Derive the Key Encryption Key (KEK) from the master password using PBKDF2
-    const masterPassword = this.configService.get('engine.masterPassword', { infer: true })
-    const kek = this.deriveKeyEncryptionKey(masterPassword)
-
-    if (!engine) {
-      // New Engine, set it up
-      engine = await this.firstTimeSetup(kek)
+    const keyringConfig = this.configService.get('keyring', { infer: true })
+
+    // We have a Raw Keyring, so we are using a MasterPassword/KEK+MasterKey for encryption
+    if (keyringConfig.masterPassword && !keyringConfig.masterAwsKmsArn) {
+      const engine = await this.encryptionRepository.getEngine(this.engineId)
+      let encryptedMasterKey = engine?.masterKey
+
+      // Derive the Key Encryption Key (KEK) from the master password using PBKDF2
+      const masterPassword = keyringConfig.masterPassword
+      const kek = this.deriveKeyEncryptionKey(masterPassword)
+
+      if (!encryptedMasterKey) {
+        // No MK yet, so create it & encrypt w/ the KEK
+        encryptedMasterKey = await this.generateMasterKey(kek)
+      }
+
+      const decryptedMasterKey = await this.decryptMasterKey(kek, Buffer.from(encryptedMasterKey, 'hex'))
+      const isolatedMasterKey = Buffer.alloc(decryptedMasterKey.length)
+      decryptedMasterKey.copy(isolatedMasterKey, 0, 0, decryptedMasterKey.length)
+
+      /* Configure the Raw AES keyring. */
+      const keyring = new RawAesKeyringNode({
+        keyName: 'armory.engine.wrapping-key',
+        keyNamespace,
+        unencryptedMasterKey: isolatedMasterKey,
+        wrappingSuite
+      })
+
+      this.keyring = keyring
+    }
+    // We have AWS KMS config so we'll use that instead as the MasterKey, which means we don't need a KEK separately
+    else if (keyringConfig.masterAwsKmsArn && !keyringConfig.masterPassword) {
+      const keyring = new KmsKeyringNode({ generatorKeyId: keyringConfig.masterAwsKmsArn })
+      this.keyring = keyring
+    } else {
+      throw new Error('Invalid Keyring Configuration found')
     }
-
-    const decryptedMasterKey = await this.decryptMasterKey(kek, Buffer.from(engine.masterKey, 'hex'))
-    this.masterKey = Buffer.alloc(decryptedMasterKey.length)
-    decryptedMasterKey.copy(decryptedMasterKey, 0, 0, decryptedMasterKey.length)
-
-    this.adminApiKey = Buffer.from(engine.adminApiKey, 'hex')
   }
 
   private getKeyEncryptionKeyring(kek: Buffer) {
@@ -71,22 +93,6 @@ export class EncryptionService implements OnApplicationBootstrap {
     return keyring
   }
 
-  private async getKeyring() {
-    if (!this.masterKey) throw new Error('Master Key not set')
-
-    /* Configure the Raw AES keyring. */
-    const keyring = new RawAesKeyringNode({
-      keyName: 'armory.engine.wrapping-key',
-      keyNamespace,
-      unencryptedMasterKey: this.masterKey,
-      wrappingSuite
-    })
-
-    // TODO: also support KMS keyring
-
-    return keyring
-  }
-
   private deriveKeyEncryptionKey(password: string): Buffer {
     // Derive the Key Encryption Key (KEK) from the master password using PBKDF2
     const kek = crypto.pbkdf2Sync(password.normalize(), this.engineId.normalize(), 1000000, 32, 'sha256')
@@ -118,7 +124,8 @@ export class EncryptionService implements OnApplicationBootstrap {
   }
 
   async encrypt(cleartext: string | Buffer): Promise<Buffer> {
-    const keyring = await this.getKeyring()
+    const keyring = this.keyring
+    if (!keyring) throw new Error('Keyring not set')
 
     const { result } = await encrypt(keyring, cleartext, {
       encryptionContext: defaultEncryptionContext
@@ -128,7 +135,8 @@ export class EncryptionService implements OnApplicationBootstrap {
   }
 
   async decrypt(ciphertext: Buffer): Promise<Buffer> {
-    const keyring = await this.getKeyring()
+    const keyring = this.keyring
+    if (!keyring) throw new Error('Keyring not set')
 
     const { plaintext, messageHeader } = await decrypt(keyring, ciphertext)
 
@@ -142,31 +150,26 @@ export class EncryptionService implements OnApplicationBootstrap {
     return plaintext
   }
 
-  private async firstTimeSetup(kek: Buffer) {
+  private async generateMasterKey(kek: Buffer): Promise<string> {
     // Generate a new Master Key (MK) with AES256
     const mk = crypto.generateKeySync('aes', { length: 256 })
     const mkBuffer = mk.export()
 
+    // Encrypt it with the Key Encryption Key (KEK) that was derived from the MP
     const encryptedMk = await this.encryptMaterKey(kek, mkBuffer)
-
-    // Generate an Admin API Key, just a random 32-byte string
-    const adminApiKeyBuffer = crypto.randomBytes(32)
+    const encryptedMkString = encryptedMk.toString('hex')
 
     // Save the Result.
-    const engine = await this.encryptionRepository.createEngine(
-      this.engineId,
-      encryptedMk.toString('hex'),
-      adminApiKeyBuffer.toString('hex') // TODO: this isn't encrypted, it should be encrypted with a CEK not with KEK
-    )
-
-    this.logger.log('Engine Initial Setup Complete')
-    // TODO: Print this to a console in a better way; may not even like this.
-    this.logger.log('Admin API Key -- DO NOT LOSE THIS', adminApiKeyBuffer.toString('hex'))
-    return engine
-  }
+    const existingEngine = await this.encryptionRepository.getEngine(this.engineId)
+    const engine = existingEngine
+      ? await this.encryptionRepository.saveMasterKey(this.engineId, encryptedMkString)
+      : await this.encryptionRepository.createEngine(this.engineId, encryptedMkString)
+
+    if (!engine?.masterKey) {
+      throw new Error('Master Key was not saved')
+    }
 
-  // Verify if a given string matches our internal Admin Api Key
-  verifyAdminApiKey(apiKey: string): boolean {
-    return apiKey === this.adminApiKey?.toString('hex')
+    this.logger.log('Engine Master Key Setup Complete')
+    return encryptedMkString
   }
 }

apps/policy-engine/src/encryption/persistence/repository/encryption.repository.ts

@@ -15,7 +15,7 @@ export class EncryptionRepository {
     })
   }
 
-  async createEngine(engineId: string, masterKey: string, adminApiKey: string) {
+  async createEngine(engineId: string, masterKey: string, adminApiKey?: string) {
     return this.prismaService.engine.create({
       data: {
         id: engineId,
@@ -24,4 +24,16 @@ export class EncryptionRepository {
       }
     })
   }
+
+  async saveMasterKey(engineId: string, masterKey: string) {
+    return this.prismaService.engine.update({
+      where: {
+        id: engineId,
+        masterKey: null // ONLY allow updating it if already null. We don't want to accidentally overwrite it!
+      },
+      data: {
+        masterKey
+      }
+    })
+  }
 }

apps/policy-engine/src/policy-engine.config.ts

@@ -13,8 +13,12 @@ const ConfigSchema = z.object({
     url: z.string().startsWith('postgresql:')
   }),
   engine: z.object({
-    id: z.string(),
-    masterPassword: z.string()
+    id: z.string()
+  }),
+  keyring: z.object({
+    type: z.enum(['awskms', 'raw']).default('raw'),
+    masterAwsKmsArn: z.string().optional(), // only if type = awskms
+    masterPassword: z.string().optional() // only if type = raw
   })
 })
 
@@ -28,7 +32,11 @@ export const load = (): Config => {
       url: process.env.POLICY_ENGINE_DATABASE_URL
     },
     engine: {
-      id: process.env.ENGINE_UID,
+      id: process.env.ENGINE_UID
+    },
+    keyring: {
+      type: process.env.KEYRING_TYPE,
+      masterAwsKmsArn: process.env.MASTER_AWS_KMS_ARN,
       masterPassword: process.env.MASTER_PASSWORD
     }
   })

apps/policy-engine/src/shared/module/persistence/schema/migrations/20240228183059_init/migration.sql → apps/policy-engine/src/shared/module/persistence/schema/migrations/20240301204146_init/migration.sql

@@ -1,8 +1,8 @@
 -- CreateTable
 CREATE TABLE "engine" (
     "id" TEXT NOT NULL,
-    "master_key" TEXT NOT NULL,
-    "admin_api_key" TEXT NOT NULL,
+    "master_key" TEXT,
+    "admin_api_key" TEXT,
 
     CONSTRAINT "engine_pkey" PRIMARY KEY ("id")
 );

apps/policy-engine/src/shared/module/persistence/schema/schema.prisma

@@ -14,8 +14,8 @@ datasource db {
 
 model Engine {
   id          String @id
-  masterKey   String @map("master_key")
-  adminApiKey String @map("admin_api_key")
+  masterKey   String? @map("master_key")
+  adminApiKey String? @map("admin_api_key")
 
   @@map("engine")
 }