Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Add @encrypted enhancer #1922

Open
wants to merge 17 commits into
base: dev
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 6 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
141 changes: 141 additions & 0 deletions packages/runtime/src/enhancements/edge/encrypted.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,141 @@
/* eslint-disable @typescript-eslint/no-explicit-any */
genu marked this conversation as resolved.
Show resolved Hide resolved
/* eslint-disable @typescript-eslint/no-unused-vars */

import { NestedWriteVisitor, enumerate, getModelFields, resolveField, type PrismaWriteActionType } from '../../cross';
import { DbClientContract } from '../../types';
import { InternalEnhancementOptions } from './create-enhancement';
import { DefaultPrismaProxyHandler, PrismaProxyActions, makeProxy } from './proxy';
import { QueryUtils } from './query-utils';

/**
* Gets an enhanced Prisma client that supports `@encrypted` attribute.
*
* @private
*/
export function withEncrypted<DbClient extends object = any>(
prisma: DbClient,
options: InternalEnhancementOptions
): DbClient {
return makeProxy(
prisma,
options.modelMeta,
(_prisma, model) => new EncryptedHandler(_prisma as DbClientContract, model, options),
'encrypted'
);
}

const encoder = new TextEncoder();
const decoder = new TextDecoder();

const getKey = async (secret: string): Promise<CryptoKey> => {
return crypto.subtle.importKey('raw', encoder.encode(secret).slice(0, 32), 'AES-GCM', false, [
'encrypt',
'decrypt',
]);
};
const encryptFunc = async (data: string, secret: string): Promise<string> => {
const key = await getKey(secret);
const iv = crypto.getRandomValues(new Uint8Array(12));

const encrypted = await crypto.subtle.encrypt(
{
name: 'AES-GCM',
iv,
},
key,
encoder.encode(data)
);

// Combine IV and encrypted data into a single array of bytes
const bytes = [...iv, ...new Uint8Array(encrypted)];

// Convert bytes to base64 string
return btoa(String.fromCharCode(...bytes));
};

const decryptFunc = async (encryptedData: string, secret: string): Promise<string> => {
const key = await getKey(secret);

// Convert base64 back to bytes
const bytes = Uint8Array.from(atob(encryptedData), (c) => c.charCodeAt(0));

// First 12 bytes are IV, rest is encrypted data
const decrypted = await crypto.subtle.decrypt(
{
name: 'AES-GCM',
iv: bytes.slice(0, 12),
},
key,
bytes.slice(12)
);

return decoder.decode(decrypted);
};

class EncryptedHandler extends DefaultPrismaProxyHandler {
private queryUtils: QueryUtils;

constructor(prisma: DbClientContract, model: string, options: InternalEnhancementOptions) {
super(prisma, model, options);

this.queryUtils = new QueryUtils(prisma, options);
}

// base override
protected async preprocessArgs(action: PrismaProxyActions, args: any) {
const actionsOfInterest: PrismaProxyActions[] = ['create', 'createMany', 'update', 'updateMany', 'upsert'];
if (args && args.data && actionsOfInterest.includes(action)) {
await this.preprocessWritePayload(this.model, action as PrismaWriteActionType, args);
}
return args;
}

// base override
protected async processResultEntity<T>(method: PrismaProxyActions, data: T): Promise<T> {
if (!data || typeof data !== 'object') {
return data;
}

for (const value of enumerate(data)) {
await this.doPostProcess(value, this.model);
}

return data;
}

private async doPostProcess(entityData: any, model: string) {
const realModel = this.queryUtils.getDelegateConcreteModel(model, entityData);

for (const field of getModelFields(entityData)) {
const fieldInfo = await resolveField(this.options.modelMeta, realModel, field);

if (!fieldInfo) {
continue;
}

const shouldDecrypt = fieldInfo.attributes?.find((attr) => attr.name === '@encrypted');
if (shouldDecrypt) {
const descryptSecret = shouldDecrypt.args.find((arg) => arg.name === 'secret')?.value as string;

entityData[field] = await decryptFunc(entityData[field], descryptSecret);
}
}
}

private async preprocessWritePayload(model: string, action: PrismaWriteActionType, args: any) {
const visitor = new NestedWriteVisitor(this.options.modelMeta, {
field: async (field, _action, data, context) => {
const encAttr = field.attributes?.find((attr) => attr.name === '@encrypted');
if (encAttr && field.type === 'String') {
// encrypt value

const secret: string = encAttr.args.find((arg) => arg.name === 'secret')?.value as string;

context.parent[field.name] = await encryptFunc(data, secret);
}
},
});

await visitor.visit(model, action, args);
}
}
15 changes: 13 additions & 2 deletions packages/runtime/src/enhancements/node/create-enhancement.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,14 @@ import { withJsonProcessor } from './json-processor';
import { Logger } from './logger';
import { withOmit } from './omit';
import { withPassword } from './password';
import { withEncrypted } from './encrypted';
import { policyProcessIncludeRelationPayload, withPolicy } from './policy';
import type { PolicyDef } from './types';

/**
* All enhancement kinds
*/
const ALL_ENHANCEMENTS: EnhancementKind[] = ['password', 'omit', 'policy', 'validation', 'delegate'];
const ALL_ENHANCEMENTS: EnhancementKind[] = ['password', 'omit', 'policy', 'validation', 'delegate', 'encrypted'];
Copy link
Contributor Author

@genu genu Dec 24, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I supposed we don't want to add the encrypted enhancement by default?

Doing so, would require the user to specify encryption options during setup, which would cause a breaking change for current users.


/**
* Options for {@link createEnhancement}
Expand Down Expand Up @@ -100,6 +101,7 @@ export function createEnhancement<DbClient extends object>(
}

const hasPassword = allFields.some((field) => field.attributes?.some((attr) => attr.name === '@password'));
const hasEncrypted = allFields.some((field) => field.attributes?.some((attr) => attr.name === '@encrypted'));
const hasOmit = allFields.some((field) => field.attributes?.some((attr) => attr.name === '@omit'));
const hasDefaultAuth = allFields.some((field) => field.defaultValueProvider);
const hasTypeDefField = allFields.some((field) => field.isTypeDef);
Expand All @@ -120,13 +122,22 @@ export function createEnhancement<DbClient extends object>(
}
}

// password enhancement must be applied prior to policy because it changes then length of the field
// password and encrypted enhancement must be applied prior to policy because it changes then length of the field
// and can break validation rules like `@length`
if (hasPassword && kinds.includes('password')) {
// @password proxy
result = withPassword(result, options);
}

if (hasEncrypted && kinds.includes('encrypted')) {
if (!options.encryption) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shall we validate the shape of options.encryption? Here or inside the EncryptedHandler constructor.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel like validating inside of the handler constructor makes better sense, and its together with the rest of the encrypted logic

throw new Error('Encryption options are required for @encrypted enhancement');
}

// @encrypted proxy
result = withEncrypted(result, options);
}

// 'policy' and 'validation' enhancements are both enabled by `withPolicy`
if (kinds.includes('policy') || kinds.includes('validation')) {
result = withPolicy(result, options, context);
Expand Down
159 changes: 159 additions & 0 deletions packages/runtime/src/enhancements/node/encrypted.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,159 @@
/* eslint-disable @typescript-eslint/no-explicit-any */
/* eslint-disable @typescript-eslint/no-unused-vars */

import {
FieldInfo,
NestedWriteVisitor,
enumerate,
getModelFields,
resolveField,
type PrismaWriteActionType,
} from '../../cross';
import { DbClientContract, CustomEncryption, SimpleEncryption } from '../../types';
import { InternalEnhancementOptions } from './create-enhancement';
import { DefaultPrismaProxyHandler, PrismaProxyActions, makeProxy } from './proxy';
import { QueryUtils } from './query-utils';

/**
* Gets an enhanced Prisma client that supports `@encrypted` attribute.
*
* @private
*/
export function withEncrypted<DbClient extends object = any>(
prisma: DbClient,
options: InternalEnhancementOptions
): DbClient {
return makeProxy(
prisma,
options.modelMeta,
(_prisma, model) => new EncryptedHandler(_prisma as DbClientContract, model, options),
'encrypted'
);
}

const encoder = new TextEncoder();
const decoder = new TextDecoder();

const getKey = async (secret: string): Promise<CryptoKey> => {
return crypto.subtle.importKey('raw', encoder.encode(secret).slice(0, 32), 'AES-GCM', false, [
'encrypt',
'decrypt',
]);
};

class EncryptedHandler extends DefaultPrismaProxyHandler {
private queryUtils: QueryUtils;

constructor(prisma: DbClientContract, model: string, options: InternalEnhancementOptions) {
super(prisma, model, options);

this.queryUtils = new QueryUtils(prisma, options);
}

private isCustomEncryption(encryption: CustomEncryption | SimpleEncryption): encryption is CustomEncryption {
return 'encrypt' in encryption && 'decrypt' in encryption;
}

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 getKey(this.options.encryption!.encryptionKey);
const iv = crypto.getRandomValues(new Uint8Array(12));

const encrypted = await crypto.subtle.encrypt(
{
name: 'AES-GCM',
iv,
},
key,
encoder.encode(data)
);

// Combine IV and encrypted data into a single array of bytes
const bytes = [...iv, ...new Uint8Array(encrypted)];

// Convert bytes to base64 string
return btoa(String.fromCharCode(...bytes));
}

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 key = await getKey(this.options.encryption!.encryptionKey);

// Convert base64 back to bytes
const bytes = Uint8Array.from(atob(data), (c) => c.charCodeAt(0));

// First 12 bytes are IV, rest is encrypted data
const decrypted = await crypto.subtle.decrypt(
{
name: 'AES-GCM',
iv: bytes.slice(0, 12),
},
key,
bytes.slice(12)
);

return decoder.decode(decrypted);
}

// base override
protected async preprocessArgs(action: PrismaProxyActions, args: any) {
const actionsOfInterest: PrismaProxyActions[] = ['create', 'createMany', 'update', 'updateMany', 'upsert'];
if (args && args.data && actionsOfInterest.includes(action)) {
await this.preprocessWritePayload(this.model, action as PrismaWriteActionType, args);
}
return args;
}

// base override
protected async processResultEntity<T>(method: PrismaProxyActions, data: T): Promise<T> {
if (!data || typeof data !== 'object') {
return data;
}

for (const value of enumerate(data)) {
await this.doPostProcess(value, this.model);
}

return data;
}

private async doPostProcess(entityData: any, model: string) {
const realModel = this.queryUtils.getDelegateConcreteModel(model, entityData);

for (const field of getModelFields(entityData)) {
const fieldInfo = await resolveField(this.options.modelMeta, realModel, field);

if (!fieldInfo) {
continue;
}

const shouldDecrypt = fieldInfo.attributes?.find((attr) => attr.name === '@encrypted');
if (shouldDecrypt) {
entityData[field] = await this.decrypt(fieldInfo, entityData[field]);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If decryption fails, should we return the original cipher text? I'm thinking this will allow easier adoption: the @encrypted attribute can be added and deployed and then a background script is run to migrate the existing plain-text data.

}
}
}

private async preprocessWritePayload(model: string, action: PrismaWriteActionType, args: any) {
const visitor = new NestedWriteVisitor(this.options.modelMeta, {
field: async (field, _action, data, context) => {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shall we add a nullish check to data here?

const encAttr = field.attributes?.find((attr) => attr.name === '@encrypted');
if (encAttr && field.type === 'String') {
// encrypt value

const secret: string = encAttr.args.find((arg) => arg.name === 'secret')?.value as string;

context.parent[field.name] = await this.encrypt(field, data);
}
},
});

await visitor.visit(model, action, args);
}
}
15 changes: 14 additions & 1 deletion packages/runtime/src/types.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
/* eslint-disable @typescript-eslint/no-explicit-any */

import type { z } from 'zod';
import { FieldInfo } from './cross';

export type PrismaPromise<T> = Promise<T> & Record<string, (args?: any) => PrismaPromise<any>>;

Expand Down Expand Up @@ -133,6 +134,11 @@ export type EnhancementOptions = {
* The `isolationLevel` option passed to `prisma.$transaction()` call for transactions initiated by ZenStack.
*/
transactionIsolationLevel?: TransactionIsolationLevel;

/**
* The encryption options for using the `encrypted` enhancement.
*/
encryption?: SimpleEncryption | CustomEncryption;
};

/**
Expand All @@ -145,7 +151,7 @@ export type EnhancementContext<User extends AuthUser = AuthUser> = {
/**
* Kinds of enhancements to `PrismaClient`
*/
export type EnhancementKind = 'password' | 'omit' | 'policy' | 'validation' | 'delegate';
export type EnhancementKind = 'password' | 'omit' | 'policy' | 'validation' | 'delegate' | 'encrypted';

/**
* Function for transforming errors.
Expand All @@ -166,3 +172,10 @@ export type ZodSchemas = {
*/
input?: Record<string, Record<string, z.ZodSchema>>;
};

export type CustomEncryption = {
encrypt: (model: string, field: FieldInfo, plain: string) => string;
decrypt: (model: string, field: FieldInfo, cipher: string) => string;
};

export type SimpleEncryption = { encryptionKey: string };
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since we don't need to represent the key in plain text anywhere, is it better to define the key as a UInt8Array instead? So that we don't need to imply any encoding form or limit it to be a utf-8 string.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we do so, we can also validate if the given key is 32-byte long.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we still need to keep it a string since the user would actually provide it as a string.

Maybe keep it as string, but ensure it is a valid UInt8Array?

@ymc9 what do you think?

Loading
Loading