fix(policy): allow dangerous raw SQL opt-in (#2502)

Author: pkudinov

packages/plugins/policy/src/options.ts

@@ -0,0 +1,8 @@
+export type PolicyPluginOptions = {
+    /**
+     * Dangerously bypasses access-policy enforcement for raw SQL queries.
+     * Raw queries remain in the current transaction, but the policy plugin will
+     * not inspect or reject them.
+     */
+    dangerouslyAllowRawSql?: boolean;
+};

packages/plugins/policy/src/plugin.ts

@@ -1,9 +1,14 @@
 import { type OnKyselyQueryArgs, type RuntimePlugin } from '@zenstackhq/orm';
 import type { SchemaDef } from '@zenstackhq/orm/schema';
+import type { PolicyPluginOptions } from './options';
 import { check } from './functions';
 import { PolicyHandler } from './policy-handler';
 
+export type { PolicyPluginOptions } from './options';
+
 export class PolicyPlugin implements RuntimePlugin<SchemaDef, {}, {}, {}> {
+    constructor(private readonly options: PolicyPluginOptions = {}) {}
+
     get id() {
         return 'policy' as const;
     }
@@ -23,7 +28,7 @@ export class PolicyPlugin implements RuntimePlugin<SchemaDef, {}, {}, {}> {
     }
 
     onKyselyQuery({ query, client, proceed }: OnKyselyQueryArgs<SchemaDef>) {
-        const handler = new PolicyHandler<SchemaDef>(client);
+        const handler = new PolicyHandler<SchemaDef>(client, this.options);
         return handler.handle(query, proceed);
     }
 }

packages/plugins/policy/src/policy-handler.ts

@@ -23,6 +23,7 @@ import {
     OperatorNode,
     ParensNode,
     PrimitiveValueListNode,
+    RawNode,
     ReferenceNode,
     ReturningNode,
     SelectAllNode,
@@ -42,6 +43,7 @@ import {
 import { match } from 'ts-pattern';
 import { ColumnCollector } from './column-collector';
 import { ExpressionTransformer } from './expression-transformer';
+import type { PolicyPluginOptions } from './options';
 import type { Policy, PolicyOperation } from './types';
 import {
     buildIsFalse,
@@ -67,7 +69,10 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
     private readonly dialect: BaseCrudDialect<Schema>;
     private readonly eb = expressionBuilder<any, any>();
 
-    constructor(private readonly client: ClientContract<Schema>) {
+    constructor(
+        private readonly client: ClientContract<Schema>,
+        private readonly options: PolicyPluginOptions = {},
+    ) {
         super();
         this.dialect = getCrudDialect(this.client.$schema, this.client.$options);
     }
@@ -76,6 +81,9 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
 
     async handle(node: RootOperationNode, proceed: ProceedKyselyQueryFunction) {
         if (!this.isCrudQueryNode(node)) {
+            if (this.options.dangerouslyAllowRawSql && RawNode.is(node as never)) {
+                return proceed(node);
+            }
             // non-CRUD queries are not allowed
             throw createRejectedByPolicyError(
                 undefined,

tests/e2e/orm/policy/raw-sql.test.ts

@@ -0,0 +1,105 @@
+import { PolicyPlugin } from '@zenstackhq/plugin-policy';
+import type { ClientContract } from '@zenstackhq/orm';
+import type { SchemaDef } from '@zenstackhq/orm/schema';
+import { createTestClient } from '@zenstackhq/testtools';
+import { sql } from 'kysely';
+import { afterEach, describe, expect, it } from 'vitest';
+
+const schema = `
+model User {
+    id String @id
+    role String
+    secrets Secret[]
+
+    @@allow('all', true)
+}
+
+model Secret {
+    id String @id
+    value String
+    ownerId String
+    owner User @relation(fields: [ownerId], references: [id])
+
+    @@allow('read', auth() != null && auth().role == 'admin')
+    @@allow('create', auth() != null && auth().role == 'admin')
+}
+`;
+
+describe('Policy raw SQL tests', () => {
+    const clients: ClientContract<SchemaDef>[] = [];
+
+    afterEach(async () => {
+        await Promise.all(clients.splice(0).map((client) => client.$disconnect()));
+    });
+
+    function ref(client: ClientContract<SchemaDef>, col: string) {
+        return client.$schema.provider.type === 'mysql' ? sql.raw(`\`${col}\``) : sql.raw(`"${col}"`);
+    }
+
+    async function createPolicyClient(options?: { dangerouslyAllowRawSql?: boolean; dbName: string }) {
+        const unsafeClient = await createTestClient(schema, {
+            dbName: options?.dbName,
+            plugins: [new PolicyPlugin({ dangerouslyAllowRawSql: options?.dangerouslyAllowRawSql })],
+        });
+        clients.push(unsafeClient);
+
+        const rawClient = unsafeClient.$unuseAll();
+        const adminClient = unsafeClient.$setAuth({ id: 'admin', role: 'admin' });
+
+        await rawClient.user.create({
+            data: {
+                id: 'admin',
+                role: 'admin',
+            },
+        });
+
+        return { adminClient };
+    }
+
+    it('keeps rejecting raw SQL by default', async () => {
+        const { adminClient } = await createPolicyClient({ dbName: 'policy_raw_sql_default' });
+
+        await expect(
+            adminClient.$transaction(async (tx) => {
+                await tx.secret.create({
+                    data: {
+                        id: 'secret-default',
+                        ownerId: 'admin',
+                        value: 'still-guarded',
+                    },
+                });
+
+                await tx.$queryRaw<{ value: string }[]>`
+                    SELECT ${ref(tx, 'Secret')}.${ref(tx, 'value')}
+                    FROM ${ref(tx, 'Secret')}
+                    WHERE ${ref(tx, 'Secret')}.${ref(tx, 'id')} = ${'secret-default'}
+                `;
+            }),
+        ).rejects.toThrow('non-CRUD queries are not allowed');
+    });
+
+    it('allows raw SQL inside a transaction when dangerous raw SQL is enabled', async () => {
+        const { adminClient } = await createPolicyClient({
+            dangerouslyAllowRawSql: true,
+            dbName: 'policy_raw_sql_dangerous',
+        });
+
+        await adminClient.$transaction(async (tx) => {
+            await tx.secret.create({
+                data: {
+                    id: 'secret-1',
+                    ownerId: 'admin',
+                    value: 'top-secret',
+                },
+            });
+
+            const rows = await tx.$queryRaw<{ value: string }[]>`
+                SELECT ${ref(tx, 'Secret')}.${ref(tx, 'value')}
+                FROM ${ref(tx, 'Secret')}
+                WHERE ${ref(tx, 'Secret')}.${ref(tx, 'id')} = ${'secret-1'}
+            `;
+
+            expect(rows).toEqual([{ value: 'top-secret' }]);
+        });
+    });
+});