fix: allow using currentModel() and currentOperation calls in policy rules (#2050)

Author: ymc9

Diff for: packages/schema/src/language-server/validator/function-invocation-validator.ts

@@ -11,6 +11,7 @@ import {
     isDataModel,
     isDataModelAttribute,
     isDataModelFieldAttribute,
+    isInvocationExpr,
     isLiteralExpr,
 } from '@zenstackhq/language/ast';
 import {
@@ -21,6 +22,7 @@ import {
     isDataModelFieldReference,
     isEnumFieldReference,
     isFromStdlib,
+    isValidationAttribute,
 } from '@zenstackhq/sdk';
 import { AstNode, streamAst, ValidationAcceptor } from 'langium';
 import { match, P } from 'ts-pattern';
@@ -70,20 +72,21 @@ export default class FunctionInvocationValidator implements AstValidator<Express
             }
 
             // validate the context allowed for the function
-            const exprContext = match(containerAttribute?.decl.$refText)
-                .with('@default', () => ExpressionContext.DefaultValue)
-                .with(P.union('@@allow', '@@deny', '@allow', '@deny'), () => ExpressionContext.AccessPolicy)
-                .with('@@validate', () => ExpressionContext.ValidationRule)
-                .with('@@index', () => ExpressionContext.Index)
-                .otherwise(() => undefined);
+            const exprContext = this.getExpressionContext(containerAttribute);
 
             // get the context allowed for the function
             const funcAllowedContext = getFunctionExpressionContext(funcDecl);
 
-            if (exprContext && !funcAllowedContext.includes(exprContext)) {
-                accept('error', `function "${funcDecl.name}" is not allowed in the current context: ${exprContext}`, {
-                    node: expr,
-                });
+            if (funcAllowedContext.length > 0 && (!exprContext || !funcAllowedContext.includes(exprContext))) {
+                accept(
+                    'error',
+                    `function "${funcDecl.name}" is not allowed in the current context${
+                        exprContext ? ': ' + exprContext : ''
+                    }`,
+                    {
+                        node: expr,
+                    }
+                );
                 return;
             }
 
@@ -121,6 +124,8 @@ export default class FunctionInvocationValidator implements AstValidator<Express
                     !isEnumFieldReference(secondArg) &&
                     // `auth()...` expression
                     !isAuthOrAuthMemberAccess(secondArg) &&
+                    // static function calls that are runtime constants: `currentModel`, `currentOperation`
+                    !this.isStaticFunctionCall(secondArg) &&
                     // array of literal/enum
                     !(
                         isArrayExpr(secondArg) &&
@@ -148,6 +153,24 @@ export default class FunctionInvocationValidator implements AstValidator<Express
         }
     }
 
+    private getExpressionContext(containerAttribute: DataModelAttribute | DataModelFieldAttribute | undefined) {
+        if (!containerAttribute) {
+            return undefined;
+        }
+        if (isValidationAttribute(containerAttribute)) {
+            return ExpressionContext.ValidationRule;
+        }
+        return match(containerAttribute?.decl.$refText)
+            .with('@default', () => ExpressionContext.DefaultValue)
+            .with(P.union('@@allow', '@@deny', '@allow', '@deny'), () => ExpressionContext.AccessPolicy)
+            .with('@@index', () => ExpressionContext.Index)
+            .otherwise(() => undefined);
+    }
+
+    private isStaticFunctionCall(expr: Expression) {
+        return isInvocationExpr(expr) && ['currentModel', 'currentOperation'].includes(expr.function.$refText);
+    }
+
     private validateArgs(funcDecl: FunctionDecl, args: Argument[], accept: ValidationAcceptor) {
         let success = true;
         for (let i = 0; i < funcDecl.params.length; i++) {

Diff for: packages/sdk/src/validation.ts

@@ -7,7 +7,7 @@ import {
     type TypeDef,
 } from './ast';
 
-function isValidationAttribute(attr: DataModelAttribute | DataModelFieldAttribute) {
+export function isValidationAttribute(attr: DataModelAttribute | DataModelFieldAttribute) {
     return attr.decl.ref?.attributes.some((attr) => attr.decl.$refText === '@@@validation');
 }
 

Diff for: tests/regression/tests/issue-1984.test.ts

@@ -0,0 +1,57 @@
+import { loadModel, loadModelWithError, loadSchema } from '@zenstackhq/testtools';
+
+describe('issue 1984', () => {
+    it('regression1', async () => {
+        const { enhance } = await loadSchema(
+            `
+            model User {
+                id Int @id @default(autoincrement())
+                access String
+
+                @@allow('all', 
+                    contains(auth().access, currentModel()) || 
+                    contains(auth().access, currentOperation()))
+            }
+            `
+        );
+
+        const db1 = enhance();
+        await expect(db1.user.create({ data: { access: 'foo' } })).toBeRejectedByPolicy();
+
+        const db2 = enhance({ id: 1, access: 'aUser' });
+        await expect(db2.user.create({ data: { access: 'aUser' } })).toResolveTruthy();
+
+        const db3 = enhance({ id: 1, access: 'do-create-read' });
+        await expect(db3.user.create({ data: { access: 'do-create-read' } })).toResolveTruthy();
+
+        const db4 = enhance({ id: 1, access: 'do-read' });
+        await expect(db4.user.create({ data: { access: 'do-read' } })).toBeRejectedByPolicy();
+    });
+
+    it('regression2', async () => {
+        await expect(
+            loadModelWithError(
+                `
+            model User {
+                id Int @id @default(autoincrement())
+                modelName String
+                @@validate(contains(modelName, currentModel()))
+            }
+            `
+            )
+        ).resolves.toContain('function "currentModel" is not allowed in the current context: ValidationRule');
+    });
+
+    it('regression3', async () => {
+        await expect(
+            loadModelWithError(
+                `
+            model User {
+                id Int @id @default(autoincrement())
+                modelName String @contains(currentModel())
+            }
+            `
+            )
+        ).resolves.toContain('function "currentModel" is not allowed in the current context: ValidationRule');
+    });
+});