fix: support check() policy function call in permission checkers (#1820)

Author: ymc9

packages/runtime/src/enhancements/node/policy/policy-utils.ts

@@ -3,6 +3,7 @@
 import deepmerge from 'deepmerge';
 import { isPlainObject } from 'is-plain-object';
 import { lowerCaseFirst } from 'lower-case-first';
+import traverse from 'traverse';
 import { upperCaseFirst } from 'upper-case-first';
 import { z, type ZodError, type ZodObject, type ZodSchema } from 'zod';
 import { fromZodError } from 'zod-validation-error';
@@ -31,7 +32,15 @@ import { getVersion } from '../../../version';
 import type { InternalEnhancementOptions } from '../create-enhancement';
 import { Logger } from '../logger';
 import { QueryUtils } from '../query-utils';
-import type { EntityChecker, ModelPolicyDef, PermissionCheckerFunc, PolicyDef, PolicyFunc } from '../types';
+import type {
+    DelegateConstraint,
+    EntityChecker,
+    ModelPolicyDef,
+    PermissionCheckerFunc,
+    PolicyDef,
+    PolicyFunc,
+    VariableConstraint,
+} from '../types';
 import { formatObject, prismaClientKnownRequestError } from '../utils';
 
 /**
@@ -667,7 +676,47 @@ export class PolicyUtil extends QueryUtils {
         }
 
         // call checker function
-        return checker({ user: this.user });
+        let result = checker({ user: this.user });
+
+        // the constraint may contain "delegate" ones that should be resolved
+        // by evaluating the corresponding checker of the delegated models
+
+        const isVariableConstraint = (value: any): value is VariableConstraint => {
+            return value && typeof value === 'object' && value.kind === 'variable';
+        };
+
+        const isDelegateConstraint = (value: any): value is DelegateConstraint => {
+            return value && typeof value === 'object' && value.kind === 'delegate';
+        };
+
+        // here we prefix the constraint variables coming from delegated checkers
+        // with the relation field name to avoid conflicts
+        const prefixConstraintVariables = (constraint: unknown, prefix: string) => {
+            return traverse(constraint).map(function (value) {
+                if (isVariableConstraint(value)) {
+                    this.update(
+                        {
+                            ...value,
+                            name: `${prefix}${value.name}`,
+                        },
+                        true
+                    );
+                }
+            });
+        };
+
+        // eslint-disable-next-line @typescript-eslint/no-this-alias
+        const that = this;
+        result = traverse(result).forEach(function (value) {
+            if (isDelegateConstraint(value)) {
+                const { model: delegateModel, relation, operation: delegateOp } = value;
+                let newValue = that.getCheckerConstraint(delegateModel, delegateOp ?? operation);
+                newValue = prefixConstraintVariables(newValue, `${relation}.`);
+                this.update(newValue, true);
+            }
+        });
+
+        return result;
     }
 
     //#endregion

packages/runtime/src/enhancements/node/types.ts

@@ -18,6 +18,11 @@ export interface CommonEnhancementOptions {
     prismaModule?: any;
 }
 
+/**
+ * CRUD operations
+ */
+export type CRUD = 'create' | 'read' | 'update' | 'delete';
+
 /**
  * Function for getting policy guard with a given context
  */
@@ -74,14 +79,26 @@ export type LogicalConstraint = {
     children: PermissionCheckerConstraint[];
 };
 
+/**
+ * Constraint delegated to another model through `check()` function call
+ * on a relation field.
+ */
+export type DelegateConstraint = {
+    kind: 'delegate';
+    model: string;
+    relation: string;
+    operation?: CRUD;
+};
+
 /**
  * Operation allowability checking constraint
  */
 export type PermissionCheckerConstraint =
     | ValueConstraint
     | VariableConstraint
     | ComparisonConstraint
-    | LogicalConstraint;
+    | LogicalConstraint
+    | DelegateConstraint;
 
 /**
  * Policy definition

packages/schema/src/plugins/enhancer/policy/constraint-transformer.ts

@@ -1,4 +1,6 @@
 import {
+    PluginError,
+    getLiteral,
     getRelationKeyPairs,
     isAuthInvocation,
     isDataModelFieldReference,
@@ -7,9 +9,11 @@ import {
 import {
     BinaryExpr,
     BooleanLiteral,
+    DataModel,
     DataModelField,
     Expression,
     ExpressionType,
+    InvocationExpr,
     LiteralExpr,
     MemberAccessExpr,
     NumberLiteral,
@@ -27,6 +31,8 @@ import {
     isUnaryExpr,
 } from '@zenstackhq/sdk/ast';
 import { P, match } from 'ts-pattern';
+import { name } from '..';
+import { isCheckInvocation } from '../../../utils/ast-utils';
 
 /**
  * Options for {@link ConstraintTransformer}.
@@ -107,6 +113,8 @@ export class ConstraintTransformer {
                 .when(isReferenceExpr, (expr) => this.transformReference(expr))
                 // top-level boolean member access expr
                 .when(isMemberAccessExpr, (expr) => this.transformMemberAccess(expr))
+                // `check()` invocation on a relation field
+                .when(isCheckInvocation, (expr) => this.transformCheckInvocation(expr as InvocationExpr))
                 .otherwise(() => this.nextVar())
         );
     }
@@ -259,6 +267,30 @@ export class ConstraintTransformer {
         return undefined;
     }
 
+    private transformCheckInvocation(expr: InvocationExpr) {
+        // transform `check()` invocation to a special "delegate" constraint kind
+        // to be evaluated at runtime
+
+        const field = expr.args[0].value as ReferenceExpr;
+        if (!field) {
+            throw new PluginError(name, 'Invalid check invocation');
+        }
+        const fieldType = field.$resolvedType?.decl as DataModel;
+
+        let operation: string | undefined = undefined;
+        if (expr.args[1]) {
+            operation = getLiteral<string>(expr.args[1].value);
+        }
+
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        const result: any = { kind: 'delegate', model: fieldType.name, relation: field.target.$refText };
+        if (operation) {
+            // operation can be explicitly specified or inferred from the context
+            result.operation = operation;
+        }
+        return JSON.stringify(result);
+    }
+
     // normalize `auth()` access undefined value to null
     private normalizeToNull(expr: string) {
         return `(${expr} ?? null)`;

tests/integration/tests/enhancements/with-policy/checker.test.ts

@@ -357,7 +357,7 @@ describe('Permission checker', () => {
         await expect(db.model.check({ operation: 'update', where: { x: 1, y: 1 } })).toResolveFalsy();
     });
 
-    it('field condition unsolvable', async () => {
+    it('field condition unsatisfiable', async () => {
         const { enhance } = await load(
             `
             model Model {
@@ -649,4 +649,115 @@ describe('Permission checker', () => {
         await expect(db.model.check({ operation: 'read', where: { value: 1 } })).toResolveTruthy();
         await expect(db.model.check({ operation: 'read', where: { value: 2 } })).toResolveTruthy();
     });
+
+    it('supports policy delegation simple', async () => {
+        const { enhance } = await load(
+            `
+            model User {
+                id Int @id @default(autoincrement())
+                foo Foo[]
+            }
+
+            model Foo {
+                id Int @id @default(autoincrement())
+                owner User @relation(fields: [ownerId], references: [id])
+                ownerId Int
+                model Model?
+                @@allow('read', auth().id == ownerId)
+                @@allow('create', auth().id != ownerId)
+                @@allow('update', auth() == owner)
+            }
+
+            model Model {
+                id Int @id @default(autoincrement())
+                foo Foo @relation(fields: [fooId], references: [id])
+                fooId Int @unique
+                @@allow('all', check(foo))
+            }
+            `,
+            { preserveTsFiles: true }
+        );
+
+        await expect(enhance().model.check({ operation: 'read' })).toResolveFalsy();
+        await expect(enhance({ id: 1 }).model.check({ operation: 'read' })).toResolveTruthy();
+
+        await expect(enhance().model.check({ operation: 'create' })).toResolveFalsy();
+        await expect(enhance({ id: 1 }).model.check({ operation: 'create' })).toResolveTruthy();
+
+        await expect(enhance().model.check({ operation: 'update' })).toResolveFalsy();
+        await expect(enhance({ id: 1 }).model.check({ operation: 'update' })).toResolveTruthy();
+
+        await expect(enhance().model.check({ operation: 'delete' })).toResolveFalsy();
+        await expect(enhance({ id: 1 }).model.check({ operation: 'delete' })).toResolveFalsy();
+    });
+
+    it('supports policy delegation explicit', async () => {
+        const { enhance } = await load(
+            `
+            model Foo {
+                id Int @id @default(autoincrement())
+                model Model?
+                @@allow('all', true)
+                @@deny('update', true)
+            }
+
+            model Model {
+                id Int @id @default(autoincrement())
+                foo Foo @relation(fields: [fooId], references: [id])
+                fooId Int @unique
+                @@allow('read', check(foo, 'update'))
+            }
+            `,
+            { preserveTsFiles: true }
+        );
+
+        await expect(enhance().model.check({ operation: 'read' })).toResolveFalsy();
+    });
+
+    it('supports policy delegation combined', async () => {
+        const { enhance } = await load(
+            `
+            model User {
+                id Int @id @default(autoincrement())
+                foo Foo[]
+            }
+
+            model Foo {
+                id Int @id @default(autoincrement())
+                owner User @relation(fields: [ownerId], references: [id])
+                ownerId Int
+                model Model?
+                @@allow('read', auth().id == ownerId)
+                @@allow('create', auth().id != ownerId)
+                @@allow('update', auth() == owner)
+            }
+
+            model Model {
+                id Int @id @default(autoincrement())
+                foo Foo @relation(fields: [fooId], references: [id])
+                fooId Int @unique
+                value Int
+                @@allow('all', check(foo) && value > 0)
+                @@deny('update', check(foo) && value == 1)
+            }
+            `,
+            { preserveTsFiles: true }
+        );
+
+        await expect(enhance().model.check({ operation: 'read' })).toResolveFalsy();
+        await expect(enhance({ id: 1 }).model.check({ operation: 'read' })).toResolveTruthy();
+        await expect(enhance({ id: 1 }).model.check({ operation: 'read', where: { value: 1 } })).toResolveTruthy();
+        await expect(enhance({ id: 1 }).model.check({ operation: 'read', where: { value: 0 } })).toResolveFalsy();
+
+        await expect(enhance().model.check({ operation: 'create' })).toResolveFalsy();
+        await expect(enhance({ id: 1 }).model.check({ operation: 'create' })).toResolveTruthy();
+        await expect(enhance({ id: 1 }).model.check({ operation: 'create', where: { value: 1 } })).toResolveTruthy();
+        await expect(enhance({ id: 1 }).model.check({ operation: 'create', where: { value: 0 } })).toResolveFalsy();
+
+        await expect(enhance().model.check({ operation: 'update' })).toResolveFalsy();
+        await expect(enhance({ id: 1 }).model.check({ operation: 'update' })).toResolveTruthy();
+        await expect(enhance({ id: 1 }).model.check({ operation: 'update', where: { value: 2 } })).toResolveTruthy();
+        await expect(enhance({ id: 1 }).model.check({ operation: 'update', where: { value: 0 } })).toResolveFalsy();
+        await expect(enhance({ id: 1 }).model.check({ operation: 'update', where: { value: 1 } })).toResolveFalsy();
+    });
 });