refactor(encryption): extract standalone encrypter/decrypter (#1945)

ymc9 · web-flow · commit 4605278f4761 · 2025-01-07T19:13:24.000+08:00
diff --git a/packages/runtime/package.json b/packages/runtime/package.json
@@ -80,6 +80,10 @@
             "types": "./zod-utils.d.ts",
             "default": "./zod-utils.js"
         },
+        "./encryption": {
+            "types": "./encryption/index.d.ts",
+            "default": "./encryption/index.js"
+        },
         "./package.json": {
             "default": "./package.json"
         }
diff --git a/packages/runtime/src/encryption/index.ts b/packages/runtime/src/encryption/index.ts
@@ -0,0 +1,67 @@
+import { _decrypt, _encrypt, ENCRYPTION_KEY_BYTES, getKeyDigest, loadKey } from './utils';
+
+/**
+ * Default encrypter
+ */
+export class Encrypter {
+    private key: CryptoKey | undefined;
+    private keyDigest: string | undefined;
+
+    constructor(private readonly encryptionKey: Uint8Array) {
+        if (encryptionKey.length !== ENCRYPTION_KEY_BYTES) {
+            throw new Error(`Encryption key must be ${ENCRYPTION_KEY_BYTES} bytes`);
+        }
+    }
+
+    /**
+     * Encrypts the given data
+     */
+    async encrypt(data: string): Promise<string> {
+        if (!this.key) {
+            this.key = await loadKey(this.encryptionKey, ['encrypt']);
+        }
+
+        if (!this.keyDigest) {
+            this.keyDigest = await getKeyDigest(this.encryptionKey);
+        }
+
+        return _encrypt(data, this.key, this.keyDigest);
+    }
+}
+
+/**
+ * Default decrypter
+ */
+export class Decrypter {
+    private keys: Array<{ key: CryptoKey; digest: string }> = [];
+
+    constructor(private readonly decryptionKeys: Uint8Array[]) {
+        if (decryptionKeys.length === 0) {
+            throw new Error('At least one decryption key must be provided');
+        }
+
+        for (const key of decryptionKeys) {
+            if (key.length !== ENCRYPTION_KEY_BYTES) {
+                throw new Error(`Decryption key must be ${ENCRYPTION_KEY_BYTES} bytes`);
+            }
+        }
+    }
+
+    /**
+     * Decrypts the given data
+     */
+    async decrypt(data: string): Promise<string> {
+        if (this.keys.length === 0) {
+            this.keys = await Promise.all(
+                this.decryptionKeys.map(async (key) => ({
+                    key: await loadKey(key, ['decrypt']),
+                    digest: await getKeyDigest(key),
+                }))
+            );
+        }
+
+        return _decrypt(data, async (digest) =>
+            this.keys.filter((entry) => entry.digest === digest).map((entry) => entry.key)
+        );
+    }
+}
diff --git a/packages/runtime/src/encryption/utils.ts b/packages/runtime/src/encryption/utils.ts
@@ -0,0 +1,96 @@
+import { z } from 'zod';
+
+export const ENCRYPTER_VERSION = 1;
+export const ENCRYPTION_KEY_BYTES = 32;
+export const IV_BYTES = 12;
+export const ALGORITHM = 'AES-GCM';
+export const KEY_DIGEST_BYTES = 8;
+
+const encoder = new TextEncoder();
+const decoder = new TextDecoder();
+
+const encryptionMetaSchema = z.object({
+    // version
+    v: z.number(),
+    // algorithm
+    a: z.string(),
+    // key digest
+    k: z.string(),
+});
+
+export async function loadKey(key: Uint8Array, keyUsages: KeyUsage[]): Promise<CryptoKey> {
+    return crypto.subtle.importKey('raw', key, ALGORITHM, false, keyUsages);
+}
+
+export async function getKeyDigest(key: Uint8Array) {
+    const rawDigest = await crypto.subtle.digest('SHA-256', key);
+    return new Uint8Array(rawDigest.slice(0, KEY_DIGEST_BYTES)).reduce(
+        (acc, byte) => acc + byte.toString(16).padStart(2, '0'),
+        ''
+    );
+}
+
+export async function _encrypt(data: string, key: CryptoKey, keyDigest: string): Promise<string> {
+    const iv = crypto.getRandomValues(new Uint8Array(IV_BYTES));
+    const encrypted = await crypto.subtle.encrypt(
+        {
+            name: ALGORITHM,
+            iv,
+        },
+        key,
+        encoder.encode(data)
+    );
+
+    // combine IV and encrypted data into a single array of bytes
+    const cipherBytes = [...iv, ...new Uint8Array(encrypted)];
+
+    // encryption metadata
+    const meta = { v: ENCRYPTER_VERSION, a: ALGORITHM, k: keyDigest };
+
+    // convert concatenated result to base64 string
+    return `${btoa(JSON.stringify(meta))}.${btoa(String.fromCharCode(...cipherBytes))}`;
+}
+
+export async function _decrypt(data: string, findKey: (digest: string) => Promise<CryptoKey[]>): Promise<string> {
+    const [metaText, cipherText] = data.split('.');
+    if (!metaText || !cipherText) {
+        throw new Error('Malformed encrypted data');
+    }
+
+    let metaObj: unknown;
+    try {
+        metaObj = JSON.parse(atob(metaText));
+    } catch (error) {
+        throw new Error('Malformed metadata');
+    }
+
+    // parse meta
+    const { a: algorithm, k: keyDigest } = encryptionMetaSchema.parse(metaObj);
+
+    // find a matching decryption key
+    const keys = await findKey(keyDigest);
+    if (keys.length === 0) {
+        throw new Error('No matching decryption key found');
+    }
+
+    // convert base64 back to bytes
+    const bytes = Uint8Array.from(atob(cipherText), (c) => c.charCodeAt(0));
+
+    // extract IV from the head
+    const iv = bytes.slice(0, IV_BYTES);
+    const cipher = bytes.slice(IV_BYTES);
+    let lastError: unknown;
+
+    for (const key of keys) {
+        let decrypted: ArrayBuffer;
+        try {
+            decrypted = await crypto.subtle.decrypt({ name: algorithm, iv }, key, cipher);
+        } catch (err) {
+            lastError = err;
+            continue;
+        }
+        return decoder.decode(decrypted);
+    }
+
+    throw lastError;
+}
diff --git a/packages/runtime/src/enhancements/node/encryption.ts b/packages/runtime/src/enhancements/node/encryption.ts
@@ -1,7 +1,6 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 /* eslint-disable @typescript-eslint/no-unused-vars */
 
-import { z } from 'zod';
 import { ACTIONS_WITH_WRITE_PAYLOAD } from '../../constants';
 import {
     FieldInfo,
@@ -11,6 +10,7 @@ import {
     resolveField,
     type PrismaWriteActionType,
 } from '../../cross';
+import { Decrypter, Encrypter } from '../../encryption';
 import { CustomEncryption, DbClientContract, SimpleEncryption } from '../../types';
 import { InternalEnhancementOptions } from './create-enhancement';
 import { Logger } from './logger';
@@ -36,27 +36,12 @@ export function withEncrypted<DbClient extends object = any>(
 
 class EncryptedHandler extends DefaultPrismaProxyHandler {
     private queryUtils: QueryUtils;
-    private encoder = new TextEncoder();
-    private decoder = new TextDecoder();
     private logger: Logger;
     private encryptionKey: CryptoKey | undefined;
     private encryptionKeyDigest: string | undefined;
     private decryptionKeys: Array<{ key: CryptoKey; digest: string }> = [];
-    private encryptionMetaSchema = z.object({
-        // version
-        v: z.number(),
-        // algorithm
-        a: z.string(),
-        // key digest
-        k: z.string(),
-    });
-
-    // constants
-    private readonly ENCRYPTION_KEY_BYTES = 32;
-    private readonly IV_BYTES = 12;
-    private readonly ALGORITHM = 'AES-GCM';
-    private readonly ENCRYPTER_VERSION = 1;
-    private readonly KEY_DIGEST_BYTES = 8;
+    private encrypter: Encrypter | undefined;
+    private decrypter: Decrypter | undefined;
 
     constructor(prisma: DbClientContract, model: string, options: InternalEnhancementOptions) {
         super(prisma, model, options);
@@ -76,138 +61,33 @@ class EncryptedHandler extends DefaultPrismaProxyHandler {
             if (!options.encryption.encryptionKey) {
                 throw this.queryUtils.unknownError('Encryption key must be provided');
             }
-            if (options.encryption.encryptionKey.length !== this.ENCRYPTION_KEY_BYTES) {
-                throw this.queryUtils.unknownError(`Encryption key must be ${this.ENCRYPTION_KEY_BYTES} bytes`);
-            }
+
+            this.encrypter = new Encrypter(options.encryption.encryptionKey);
+            this.decrypter = new Decrypter([
+                options.encryption.encryptionKey,
+                ...(options.encryption.decryptionKeys || []),
+            ]);
         }
     }
 
     private isCustomEncryption(encryption: CustomEncryption | SimpleEncryption): encryption is CustomEncryption {
         return 'encrypt' in encryption && 'decrypt' in encryption;
     }
 
-    private async loadKey(key: Uint8Array, keyUsages: KeyUsage[]): Promise<CryptoKey> {
-        return crypto.subtle.importKey('raw', key, this.ALGORITHM, false, keyUsages);
-    }
-
-    private async computeKeyDigest(key: Uint8Array) {
-        const rawDigest = await crypto.subtle.digest('SHA-256', key);
-        return new Uint8Array(rawDigest.slice(0, this.KEY_DIGEST_BYTES)).reduce(
-            (acc, byte) => acc + byte.toString(16).padStart(2, '0'),
-            ''
-        );
-    }
-
-    private async getEncryptionKey(): Promise<CryptoKey> {
-        if (this.isCustomEncryption(this.options.encryption!)) {
-            throw new Error('Unexpected custom encryption settings');
-        }
-        if (!this.encryptionKey) {
-            this.encryptionKey = await this.loadKey(this.options.encryption!.encryptionKey, ['encrypt', 'decrypt']);
-        }
-        return this.encryptionKey;
-    }
-
-    private async getEncryptionKeyDigest() {
-        if (this.isCustomEncryption(this.options.encryption!)) {
-            throw new Error('Unexpected custom encryption settings');
-        }
-        if (!this.encryptionKeyDigest) {
-            this.encryptionKeyDigest = await this.computeKeyDigest(this.options.encryption!.encryptionKey);
-        }
-        return this.encryptionKeyDigest;
-    }
-
-    private async findDecryptionKeys(keyDigest: string): Promise<CryptoKey[]> {
-        if (this.isCustomEncryption(this.options.encryption!)) {
-            throw new Error('Unexpected custom encryption settings');
-        }
-
-        if (this.decryptionKeys.length === 0) {
-            const keys = [this.options.encryption!.encryptionKey, ...(this.options.encryption!.decryptionKeys || [])];
-            this.decryptionKeys = await Promise.all(
-                keys.map(async (key) => ({
-                    key: await this.loadKey(key, ['decrypt']),
-                    digest: await this.computeKeyDigest(key),
-                }))
-            );
-        }
-
-        return this.decryptionKeys.filter((entry) => entry.digest === keyDigest).map((entry) => entry.key);
-    }
-
     private async encrypt(field: FieldInfo, data: string): Promise<string> {
         if (this.isCustomEncryption(this.options.encryption!)) {
             return this.options.encryption.encrypt(this.model, field, data);
         }
 
-        const key = await this.getEncryptionKey();
-        const iv = crypto.getRandomValues(new Uint8Array(this.IV_BYTES));
-        const encrypted = await crypto.subtle.encrypt(
-            {
-                name: this.ALGORITHM,
-                iv,
-            },
-            key,
-            this.encoder.encode(data)
-        );
-
-        // combine IV and encrypted data into a single array of bytes
-        const cipherBytes = [...iv, ...new Uint8Array(encrypted)];
-
-        // encryption metadata
-        const meta = { v: this.ENCRYPTER_VERSION, a: this.ALGORITHM, k: await this.getEncryptionKeyDigest() };
-
-        // convert concatenated result to base64 string
-        return `${btoa(JSON.stringify(meta))}.${btoa(String.fromCharCode(...cipherBytes))}`;
+        return this.encrypter!.encrypt(data);
     }
 
     private async decrypt(field: FieldInfo, data: string): Promise<string> {
         if (this.isCustomEncryption(this.options.encryption!)) {
             return this.options.encryption.decrypt(this.model, field, data);
         }
 
-        const [metaText, cipherText] = data.split('.');
-        if (!metaText || !cipherText) {
-            throw new Error('Malformed encrypted data');
-        }
-
-        let metaObj: unknown;
-        try {
-            metaObj = JSON.parse(atob(metaText));
-        } catch (error) {
-            throw new Error('Malformed metadata');
-        }
-
-        // parse meta
-        const { a: algorithm, k: keyDigest } = this.encryptionMetaSchema.parse(metaObj);
-
-        // find a matching decryption key
-        const keys = await this.findDecryptionKeys(keyDigest);
-        if (keys.length === 0) {
-            throw new Error('No matching decryption key found');
-        }
-
-        // convert base64 back to bytes
-        const bytes = Uint8Array.from(atob(cipherText), (c) => c.charCodeAt(0));
-
-        // extract IV from the head
-        const iv = bytes.slice(0, this.IV_BYTES);
-        const cipher = bytes.slice(this.IV_BYTES);
-        let lastError: unknown;
-
-        for (const key of keys) {
-            let decrypted: ArrayBuffer;
-            try {
-                decrypted = await crypto.subtle.decrypt({ name: algorithm, iv }, key, cipher);
-            } catch (err) {
-                lastError = err;
-                continue;
-            }
-            return this.decoder.decode(decrypted);
-        }
-
-        throw lastError;
+        return this.decrypter!.decrypt(data);
     }
 
     // base override
