Skip to content

Commit 6bff7f4

Browse files
genugenu
committed
test: Add test
1 parent f8ee204 commit 6bff7f4

File tree

3 files changed

+190
-7
lines changed

3 files changed

+190
-7
lines changed

‎packages/runtime/src/enhancements/edge/encrypted.ts‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
/* eslint-disable @typescript-eslint/no-explicit-any */
2+
/* eslint-disable @typescript-eslint/no-unused-vars */
3+
4+
import { NestedWriteVisitor, enumerate, getModelFields, resolveField, type PrismaWriteActionType } from '../../cross';
5+
import { DbClientContract } from '../../types';
6+
import { InternalEnhancementOptions } from './create-enhancement';
7+
import { DefaultPrismaProxyHandler, PrismaProxyActions, makeProxy } from './proxy';
8+
import { QueryUtils } from './query-utils';
9+
10+
/**
11+
* Gets an enhanced Prisma client that supports `@encrypted` attribute.
12+
*
13+
* @private
14+
*/
15+
export function withEncrypted<DbClient extends object = any>(
16+
prisma: DbClient,
17+
options: InternalEnhancementOptions
18+
): DbClient {
19+
return makeProxy(
20+
prisma,
21+
options.modelMeta,
22+
(_prisma, model) => new EncryptedHandler(_prisma as DbClientContract, model, options),
23+
'encrypted'
24+
);
25+
}
26+
27+
const encoder = new TextEncoder();
28+
const decoder = new TextDecoder();
29+
30+
const getKey = async (secret: string): Promise<CryptoKey> => {
31+
return crypto.subtle.importKey('raw', encoder.encode(secret).slice(0, 32), 'AES-GCM', false, [
32+
'encrypt',
33+
'decrypt',
34+
]);
35+
};
36+
const encryptFunc = async (data: string, secret: string): Promise<string> => {
37+
const key = await getKey(secret);
38+
const iv = crypto.getRandomValues(new Uint8Array(12));
39+
40+
const encrypted = await crypto.subtle.encrypt(
41+
{
42+
name: 'AES-GCM',
43+
iv,
44+
},
45+
key,
46+
encoder.encode(data)
47+
);
48+
49+
// Combine IV and encrypted data into a single array of bytes
50+
const bytes = [...iv, ...new Uint8Array(encrypted)];
51+
52+
// Convert bytes to base64 string
53+
return btoa(String.fromCharCode(...bytes));
54+
};
55+
56+
const decryptFunc = async (encryptedData: string, secret: string): Promise<string> => {
57+
const key = await getKey(secret);
58+
59+
// Convert base64 back to bytes
60+
const bytes = Uint8Array.from(atob(encryptedData), (c) => c.charCodeAt(0));
61+
62+
// First 12 bytes are IV, rest is encrypted data
63+
const decrypted = await crypto.subtle.decrypt(
64+
{
65+
name: 'AES-GCM',
66+
iv: bytes.slice(0, 12),
67+
},
68+
key,
69+
bytes.slice(12)
70+
);
71+
72+
return decoder.decode(decrypted);
73+
};
74+
75+
class EncryptedHandler extends DefaultPrismaProxyHandler {
76+
private queryUtils: QueryUtils;
77+
78+
constructor(prisma: DbClientContract, model: string, options: InternalEnhancementOptions) {
79+
super(prisma, model, options);
80+
81+
this.queryUtils = new QueryUtils(prisma, options);
82+
}
83+
84+
// base override
85+
protected async preprocessArgs(action: PrismaProxyActions, args: any) {
86+
const actionsOfInterest: PrismaProxyActions[] = ['create', 'createMany', 'update', 'updateMany', 'upsert'];
87+
if (args && args.data && actionsOfInterest.includes(action)) {
88+
await this.preprocessWritePayload(this.model, action as PrismaWriteActionType, args);
89+
}
90+
return args;
91+
}
92+
93+
// base override
94+
protected async processResultEntity<T>(method: PrismaProxyActions, data: T): Promise<T> {
95+
if (!data || typeof data !== 'object') {
96+
return data;
97+
}
98+
99+
for (const value of enumerate(data)) {
100+
await this.doPostProcess(value, this.model);
101+
}
102+
103+
return data;
104+
}
105+
106+
private async doPostProcess(entityData: any, model: string) {
107+
const realModel = this.queryUtils.getDelegateConcreteModel(model, entityData);
108+
109+
for (const field of getModelFields(entityData)) {
110+
const fieldInfo = await resolveField(this.options.modelMeta, realModel, field);
111+
112+
if (!fieldInfo) {
113+
continue;
114+
}
115+
116+
const shouldDecrypt = fieldInfo.attributes?.find((attr) => attr.name === '@encrypted');
117+
if (shouldDecrypt) {
118+
const descryptSecret = shouldDecrypt.args.find((arg) => arg.name === 'secret')?.value as string;
119+
120+
entityData[field] = await decryptFunc(entityData[field], descryptSecret);
121+
}
122+
}
123+
}
124+
125+
private async preprocessWritePayload(model: string, action: PrismaWriteActionType, args: any) {
126+
const visitor = new NestedWriteVisitor(this.options.modelMeta, {
127+
field: async (field, _action, data, context) => {
128+
const encAttr = field.attributes?.find((attr) => attr.name === '@encrypted');
129+
if (encAttr && field.type === 'String') {
130+
// encrypt value
131+
132+
const secret: string = encAttr.args.find((arg) => arg.name === 'secret')?.value as string;
133+
134+
context.parent[field.name] = await encryptFunc(data, secret);
135+
}
136+
},
137+
});
138+
139+
await visitor.visit(model, action, args);
140+
}
141+
}

‎packages/runtime/src/enhancements/node/encrypted.ts‎

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
/* eslint-disable @typescript-eslint/no-explicit-any */
22
/* eslint-disable @typescript-eslint/no-unused-vars */
33

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

910
/**
1011
* Gets an enhanced Prisma client that supports `@encrypted` attribute.
@@ -72,8 +73,12 @@ const decryptFunc = async (encryptedData: string, secret: string): Promise<strin
7273
};
7374

7475
class EncryptedHandler extends DefaultPrismaProxyHandler {
76+
private queryUtils: QueryUtils;
77+
7578
constructor(prisma: DbClientContract, model: string, options: InternalEnhancementOptions) {
7679
super(prisma, model, options);
80+
81+
this.queryUtils = new QueryUtils(prisma, options);
7782
}
7883

7984
// base override
@@ -86,8 +91,35 @@ class EncryptedHandler extends DefaultPrismaProxyHandler {
8691
}
8792

8893
// base override
89-
protected async processResultEntity(action: PrismaProxyActions, args: any) {
90-
return args;
94+
protected async processResultEntity<T>(method: PrismaProxyActions, data: T): Promise<T> {
95+
if (!data || typeof data !== 'object') {
96+
return data;
97+
}
98+
99+
for (const value of enumerate(data)) {
100+
await this.doPostProcess(value, this.model);
101+
}
102+
103+
return data;
104+
}
105+
106+
private async doPostProcess(entityData: any, model: string) {
107+
const realModel = this.queryUtils.getDelegateConcreteModel(model, entityData);
108+
109+
for (const field of getModelFields(entityData)) {
110+
const fieldInfo = await resolveField(this.options.modelMeta, realModel, field);
111+
112+
if (!fieldInfo) {
113+
continue;
114+
}
115+
116+
const shouldDecrypt = fieldInfo.attributes?.find((attr) => attr.name === '@encrypted');
117+
if (shouldDecrypt) {
118+
const descryptSecret = shouldDecrypt.args.find((arg) => arg.name === 'secret')?.value as string;
119+
120+
entityData[field] = await decryptFunc(entityData[field], descryptSecret);
121+
}
122+
}
91123
}
92124

93125
private async preprocessWritePayload(model: string, action: PrismaWriteActionType, args: any) {
@@ -97,7 +129,7 @@ class EncryptedHandler extends DefaultPrismaProxyHandler {
97129
if (encAttr && field.type === 'String') {
98130
// encrypt value
99131

100-
let secret: string = encAttr.args.find((arg) => arg.name === 'secret')?.value as string;
132+
const secret: string = encAttr.args.find((arg) => arg.name === 'secret')?.value as string;
101133

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

‎tests/integration/tests/enhancements/with-encrypted/with-encrypted.test.ts‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,22 +13,32 @@ describe('Encrypted test', () => {
1313
});
1414

1515
it('encrypted tests', async () => {
16+
const ENCRYPTION_KEY = 'c558Gq0YQK2QcqtkMF9BGXHCQn4dMF8w';
17+
1618
const { enhance } = await loadSchema(`
1719
model User {
1820
id String @id @default(cuid())
19-
encrypted_value String @encrypted(saltLength: 16)
21+
encrypted_value String @encrypted(secret: "${ENCRYPTION_KEY}")
2022
2123
@@allow('all', true)
2224
}`);
2325

2426
const db = enhance();
25-
const r = await db.user.create({
27+
28+
const create = await db.user.create({
2629
data: {
2730
id: '1',
2831
encrypted_value: 'abc123',
2932
},
3033
});
3134

32-
expect(r.encrypted_value).toBe('abc123');
35+
const read = await db.user.findUnique({
36+
where: {
37+
id: '1',
38+
},
39+
});
40+
41+
expect(create.encrypted_value).toBe('abc123');
42+
expect(read.encrypted_value).toBe('abc123');
3343
});
3444
});

0 commit comments

Comments
 (0)