Feat: Allow custom cipher and decipher functions

src/encryption.test.ts

@@ -1,8 +1,15 @@
 import { formatKey } from '@47ng/cloak/dist/key'
-import { configureKeys } from './encryption'
+import type { DMMFModels } from './dmmf'
+import { configureKeys, decryptOnRead, encryptOnWrite } from './encryption'
 import { errors } from './errors'
+import type { MiddlewareParams } from './types'
 
-const TEST_KEY = 'k1.aesgcm256.DbQoar8ZLuUsOHZNyrnjlskInHDYlzF3q6y1KGM7DUM='
+const ENCRYPTION_TEST_KEY =
+  'k1.aesgcm256.DbQoar8ZLuUsOHZNyrnjlskInHDYlzF3q6y1KGM7DUM='
+const DECRYPTION_TEST_KEYS = [
+  'k1.aesgcm256.4BNYdJnjOQJP2adq9cGM9kb4dZxDujUs6aPS0VeRtAM=',
+  'k1.aesgcm256.El9unG7WBAVRQdATOyMggE3XrLV2ZjTGKdajfmIeBPs='
+]
 
 describe('encryption', () => {
   describe('configureKeys', () => {
@@ -13,47 +20,143 @@ describe('encryption', () => {
 
     test('Providing encryptionKey directly', () => {
       const { encryptionKey } = configureKeys({
-        encryptionKey: TEST_KEY
+        encryptionKey: ENCRYPTION_TEST_KEY
       })
-      expect(formatKey(encryptionKey.raw as Uint8Array)).toEqual(TEST_KEY)
+
+      expect(formatKey(encryptionKey.raw as Uint8Array)).toEqual(
+        ENCRYPTION_TEST_KEY
+      )
     })
 
     test('Providing encryptionKey via the environment', () => {
-      process.env.PRISMA_FIELD_ENCRYPTION_KEY = TEST_KEY
+      process.env.PRISMA_FIELD_ENCRYPTION_KEY = ENCRYPTION_TEST_KEY
       const { encryptionKey } = configureKeys({})
-      expect(formatKey(encryptionKey.raw as Uint8Array)).toEqual(TEST_KEY)
+      expect(formatKey(encryptionKey.raw as Uint8Array)).toEqual(
+        ENCRYPTION_TEST_KEY
+      )
       process.env.PRISMA_FIELD_ENCRYPTION_KEY = undefined
     })
 
     test('Encryption key is in the keychain', () => {
       const { encryptionKey, keychain } = configureKeys({
-        encryptionKey: TEST_KEY
+        encryptionKey: ENCRYPTION_TEST_KEY
       })
       expect(keychain[encryptionKey.fingerprint].key).toEqual(encryptionKey)
     })
 
     test('Loading decryption keys directly', () => {
       const { keychain } = configureKeys({
-        encryptionKey: TEST_KEY,
-        decryptionKeys: [
-          'k1.aesgcm256.4BNYdJnjOQJP2adq9cGM9kb4dZxDujUs6aPS0VeRtAM=',
-          'k1.aesgcm256.El9unG7WBAVRQdATOyMggE3XrLV2ZjTGKdajfmIeBPs='
-        ]
+        encryptionKey: ENCRYPTION_TEST_KEY,
+        decryptionKeys: DECRYPTION_TEST_KEYS
       })
       expect(Object.values(keychain).length).toEqual(3)
     })
 
     test('Loading decryption keys via the environment', () => {
-      process.env.PRISMA_FIELD_DECRYPTION_KEYS = [
-        'k1.aesgcm256.4BNYdJnjOQJP2adq9cGM9kb4dZxDujUs6aPS0VeRtAM=',
-        'k1.aesgcm256.El9unG7WBAVRQdATOyMggE3XrLV2ZjTGKdajfmIeBPs='
-      ].join(',')
+      process.env.PRISMA_FIELD_DECRYPTION_KEYS = DECRYPTION_TEST_KEYS.join(',')
 
       const { keychain } = configureKeys({
-        encryptionKey: TEST_KEY
+        encryptionKey: ENCRYPTION_TEST_KEY
       })
       expect(Object.values(keychain).length).toEqual(3)
       process.env.PRISMA_FIELD_DECRYPTION_KEYS = undefined
     })
   })
+
+  describe('encryptOnWrite', () => {
+    test('Should call custom cypher encrypt function', async () => {
+      const encryptFunction = jest.fn(
+        (decrypted: string) => `fake-encryption-${decrypted}`
+      )
+
+      const name = 'value'
+      const params: MiddlewareParams = {
+        model: 'User',
+        action: 'create',
+        args: { data: { name } },
+        runInTransaction: true,
+        dataPath: ['any']
+      }
+
+      const dmmfModels: DMMFModels = {
+        User: {
+          connections: {
+            'fake-connection': {
+              modelName: 'User',
+              isList: false
+            }
+          },
+          fields: {
+            name: {
+              encrypt: true,
+              strictDecryption: false
+            }
+          }
+        }
+      }
+
+      const keys = configureKeys({
+        encryptionKey: ENCRYPTION_TEST_KEY,
+        decryptionKeys: DECRYPTION_TEST_KEYS
+      })
+
+      encryptOnWrite(params, keys, dmmfModels, 'User.create', encryptFunction)
+
+      expect(encryptFunction).toBeCalledTimes(1)
+      expect(encryptFunction).toBeCalledWith(name)
+    })
+  })
+
+  describe('decryptOnRead', () => {
+    test('Should call custom cypher decrypt function', async () => {
+      const decryptFunction = jest.fn(
+        (encrypted: string) => `fake-decryption-${encrypted}`
+      )
+
+      const params: MiddlewareParams = {
+        model: 'User',
+        action: 'findFirst',
+        args: { where: { name: 'value' } },
+        runInTransaction: true,
+        dataPath: ['any']
+      }
+
+      const dmmfModels: DMMFModels = {
+        User: {
+          connections: {
+            'fake-connection': {
+              modelName: 'User',
+              isList: false
+            }
+          },
+          fields: {
+            name: {
+              encrypt: true,
+              strictDecryption: false
+            }
+          }
+        }
+      }
+
+      const keys = configureKeys({
+        encryptionKey: ENCRYPTION_TEST_KEY,
+        decryptionKeys: DECRYPTION_TEST_KEYS
+      })
+
+      const encryptedName = 'a1b2c3d4e5d6'
+      const result = { name: encryptedName }
+
+      decryptOnRead(
+        params,
+        result,
+        keys,
+        dmmfModels,
+        'User.findFirst',
+        decryptFunction
+      )
+
+      expect(decryptFunction).toBeCalledTimes(1)
+      expect(decryptFunction).toBeCalledWith(encryptedName)
+    })
+  })
 })

src/encryption.ts

@@ -12,15 +12,20 @@ import produce, { Draft } from 'immer'
 import objectPath from 'object-path'
 import type { DMMFModels } from './dmmf'
 import { errors, warnings } from './errors'
-import type { Configuration, MiddlewareParams } from './types'
+import type { MiddlewareParams, EncryptionFn, DecryptionFn } from './types'
 import { visitInputTargetFields, visitOutputTargetFields } from './visitor'
 
 export interface KeysConfiguration {
   encryptionKey: ParsedCloakKey
   keychain: CloakKeychain
 }
 
-export function configureKeys(config: Configuration): KeysConfiguration {
+export interface ConfigureKeysParams {
+  encryptionKey?: string
+  decryptionKeys?: string[]
+}
+
+export function configureKeys(config: ConfigureKeysParams): KeysConfiguration {
   const encryptionKey =
     config.encryptionKey || process.env.PRISMA_FIELD_ENCRYPTION_KEY
 
@@ -63,7 +68,8 @@ export function encryptOnWrite(
   params: MiddlewareParams,
   keys: KeysConfiguration,
   models: DMMFModels,
-  operation: string
+  operation: string,
+  encryptFn?: EncryptionFn
 ) {
   if (!writeOperations.includes(params.action)) {
     return params // No input data to encrypt
@@ -89,7 +95,11 @@ export function encryptOnWrite(
           console.warn(warnings.whereClause(operation, path))
         }
         try {
-          const cipherText = encryptStringSync(clearText, keys.encryptionKey)
+          const cipherText =
+            encryptFn !== undefined
+              ? encryptFn(clearText)
+              : encryptStringSync(clearText, keys.encryptionKey)
+
           objectPath.set(draft.args, path, cipherText)
         } catch (error) {
           encryptionErrors.push(
@@ -110,7 +120,8 @@ export function decryptOnRead(
   result: any,
   keys: KeysConfiguration,
   models: DMMFModels,
-  operation: string
+  operation: string,
+  decryptFn?: DecryptionFn
 ) {
   // Analyse the query to see if there's anything to decrypt.
   const model = models[params.model!]
@@ -137,11 +148,18 @@ export function decryptOnRead(
       field
     }) {
       try {
-        if (!cloakedStringRegex.test(cipherText)) {
+        if (!decryptFn && !cloakedStringRegex.test(cipherText)) {
           return
         }
-        const decryptionKey = findKeyForMessage(cipherText, keys.keychain)
-        const clearText = decryptStringSync(cipherText, decryptionKey)
+
+        const clearText =
+          decryptFn !== undefined
+            ? decryptFn(cipherText)
+            : decryptStringSync(
+                cipherText,
+                findKeyForMessage(cipherText, keys.keychain)
+              )
+
         objectPath.set(result, path, clearText)
       } catch (error) {
         const message = errors.fieldDecryptionError(model, field, path, error)

src/index.ts

@@ -1,13 +1,23 @@
 import { analyseDMMF } from './dmmf'
-import { configureKeys, decryptOnRead, encryptOnWrite } from './encryption'
+import {
+  configureKeys,
+  decryptOnRead,
+  encryptOnWrite,
+  ConfigureKeysParams
+} from './encryption'
 import type { Configuration, Middleware, MiddlewareParams } from './types'
 
 export function fieldEncryptionMiddleware(
   config: Configuration = {}
 ): Middleware {
+  const configureKeysParams: ConfigureKeysParams = {
+    encryptionKey: config.encryptionKey,
+    decryptionKeys: config.decryptionKeys
+  }
+
   // This will throw if the encryption key is missing
   // or if anything is invalid.
-  const keys = configureKeys(config)
+  const keys = configureKeys(configureKeysParams)
   const models = analyseDMMF()
 
   return async function fieldEncryptionMiddleware(
@@ -21,9 +31,25 @@ export function fieldEncryptionMiddleware(
     const operation = `${params.model}.${params.action}`
     // Params are mutated in-place for modifications to occur.
     // See https://github.com/prisma/prisma/issues/9522
-    const encryptedParams = encryptOnWrite(params, keys, models, operation)
+    const encryptedParams = encryptOnWrite(
+      params,
+      keys,
+      models,
+      operation,
+      config.cipher?.encrypt
+    )
+
     let result = await next(encryptedParams)
-    decryptOnRead(encryptedParams, result, keys, models, operation)
+
+    decryptOnRead(
+      encryptedParams,
+      result,
+      keys,
+      models,
+      operation,
+      config.cipher?.decrypt
+    )
+
     return result
   }
 }

src/types.ts

@@ -8,9 +8,18 @@ export type DMMF = typeof Prisma.dmmf
 
 // Internal types --
 
+export type EncryptionFn = (clearText: string) => string
+export type DecryptionFn = (cipherText: string) => string
+
+export type CipherFunctions = {
+  encrypt: EncryptionFn
+  decrypt: DecryptionFn
+}
+
 export interface Configuration {
   encryptionKey?: string
   decryptionKeys?: string[]
+  cipher?: CipherFunctions
 }
 
 export interface FieldConfiguration {