fix: validating currentModel and currentOperation properly

packages/language/src/validators/function-invocation-validator.ts

@@ -66,12 +66,7 @@ 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);
@@ -103,6 +98,24 @@ export default class FunctionInvocationValidator implements AstValidator<Express
         }
     }
 
+    private getExpressionContext(containerAttribute: DataModelAttribute | DataFieldAttribute | undefined) {
+        if (!containerAttribute) {
+            return undefined;
+        }
+        if (this.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 isValidationAttribute(attr: DataModelAttribute | DataFieldAttribute) {
+        return !!attr.decl.ref?.attributes.some((attr) => attr.decl.$refText === '@@@validation');
+    }
+
     private validateArgs(funcDecl: FunctionDecl, args: Argument[], accept: ValidationAcceptor) {
         let success = true;
         for (let i = 0; i < funcDecl.params.length; i++) {

tests/regression/test/v2-migrated/issue-1955.test.ts

@@ -0,0 +1,100 @@
+import { createPolicyTestClient } from '@zenstackhq/testtools';
+import { describe, expect, it } from 'vitest';
+
+describe('Regression for issue 1955', () => {
+    it('simple policy', async () => {
+        const db = await createPolicyTestClient(
+            `
+            model Post {
+                id Int @id @default(autoincrement())
+                name String
+                expections String[]
+                @@allow('all', true)
+            }
+            `,
+            { provider: 'postgresql' },
+        );
+
+        await expect(
+            db.post.createManyAndReturn({
+                data: [
+                    {
+                        name: 'bla',
+                    },
+                    {
+                        name: 'blu',
+                    },
+                ],
+            }),
+        ).resolves.toEqual(
+            expect.arrayContaining([
+                expect.objectContaining({ name: 'bla' }),
+                expect.objectContaining({ name: 'blu' }),
+            ]),
+        );
+
+        await expect(
+            db.post.updateManyAndReturn({
+                data: { name: 'foo' },
+            }),
+        ).resolves.toEqual(
+            expect.arrayContaining([
+                expect.objectContaining({ name: 'foo' }),
+                expect.objectContaining({ name: 'foo' }),
+            ]),
+        );
+    });
+
+    it('complex policy', async () => {
+        const db = await createPolicyTestClient(
+            `
+            model Post {
+                id Int @id @default(autoincrement())
+                name String
+                expections String[]
+                comments Comment[]
+
+                @@allow('create', true)
+                @@allow('read,update', comments^[private])
+            }
+
+            model Comment {
+                id Int @id @default(autoincrement())
+                private Boolean @default(false)
+                postId Int
+                post Post @relation(fields: [postId], references: [id])
+            }
+            `,
+            { provider: 'postgresql' },
+        );
+
+        await expect(
+            db.post.createManyAndReturn({
+                data: [
+                    {
+                        name: 'bla',
+                    },
+                    {
+                        name: 'blu',
+                    },
+                ],
+            }),
+        ).resolves.toEqual(
+            expect.arrayContaining([
+                expect.objectContaining({ name: 'bla' }),
+                expect.objectContaining({ name: 'blu' }),
+            ]),
+        );
+
+        await expect(
+            db.post.updateManyAndReturn({
+                data: { name: 'foo' },
+            }),
+        ).resolves.toEqual(
+            expect.arrayContaining([
+                expect.objectContaining({ name: 'foo' }),
+                expect.objectContaining({ name: 'foo' }),
+            ]),
+        );
+    });
+});

tests/regression/test/v2-migrated/issue-1964.test.ts

@@ -0,0 +1,123 @@
+import { createPolicyTestClient } from '@zenstackhq/testtools';
+import { describe, expect, it } from 'vitest';
+
+describe('Regression for issue 1964', () => {
+    it('regression1', async () => {
+        const db = await createPolicyTestClient(
+            `
+model User {
+  id    Int    @id
+  orgId String
+}
+
+model Author {
+  id    Int    @id @default(autoincrement())
+  orgId String
+  name  String
+  posts Post[]
+
+  @@unique([orgId, name])
+  @@allow('all', auth().orgId == orgId)
+}
+
+model Post {
+  id       Int     @id @default(autoincrement())
+  orgId    String
+  title    String
+  author   Author  @relation(fields: [authorId], references: [id])
+  authorId Int
+
+  @@allow('all', auth().orgId == orgId)
+}
+            `,
+        );
+
+        const authDb = db.$setAuth({ id: 1, orgId: 'org' });
+
+        const newauthor = await authDb.author.create({
+            data: {
+                name: `Foo ${Date.now()}`,
+                orgId: 'org',
+                posts: {
+                    createMany: { data: [{ title: 'Hello', orgId: 'org' }] },
+                },
+            },
+            include: { posts: true },
+        });
+
+        await expect(
+            authDb.author.update({
+                where: { orgId_name: { orgId: 'org', name: newauthor.name } },
+                data: {
+                    name: `Bar ${Date.now()}`,
+                    posts: { deleteMany: { id: { equals: newauthor.posts[0].id } } },
+                },
+            }),
+        ).toResolveTruthy();
+    });
+
+    it('regression2', async () => {
+        const db = await createPolicyTestClient(
+            `
+model User {
+  id    Int    @id @default(autoincrement())
+  slug  String @unique
+  profile Profile?
+  @@allow('all', true)
+}
+
+model Profile {
+  id    Int    @id @default(autoincrement())
+  slug  String @unique
+  name  String
+  addresses Address[]
+  userId Int? @unique
+  user User? @relation(fields: [userId], references: [id])
+  @@allow('all', true)
+}
+
+model Address {
+  id    Int    @id @default(autoincrement())
+  profileId Int @unique
+  profile Profile @relation(fields: [profileId], references: [id])
+  city String
+  @@allow('all', true)
+}
+            `,
+        );
+
+        const authDb = db.$setAuth({ id: 1, orgId: 'org' });
+
+        await authDb.user.create({
+            data: {
+                slug: `user1`,
+                profile: {
+                    create: {
+                        name: `My Profile`,
+                        slug: 'profile1',
+                        addresses: {
+                            create: { id: 1, city: 'City' },
+                        },
+                    },
+                },
+            },
+        });
+
+        await expect(
+            authDb.user.update({
+                where: { slug: 'user1' },
+                data: {
+                    profile: {
+                        update: {
+                            addresses: {
+                                deleteMany: { id: { equals: 1 } },
+                            },
+                        },
+                    },
+                },
+            }),
+        ).toResolveTruthy();
+
+        await expect(authDb.address.count()).resolves.toEqual(0);
+    });
+});

tests/regression/test/v2-migrated/issue-1978.test.ts

@@ -0,0 +1,39 @@
+import { createPolicyTestClient } from '@zenstackhq/testtools';
+import { expect, it } from 'vitest';
+
+// TODO: field-level policy support
+it.skip('regression', async () => {
+    const db = await createPolicyTestClient(
+        `
+model User {
+    id Int @id
+    posts Post[]
+    secret String @allow('read', posts?[published])
+    @@allow('all', true)
+}
+
+model Post {
+    id Int @id
+    author User @relation(fields: [authorId], references: [id])
+    authorId Int
+    published Boolean @default(false)
+    @@allow('all', true)
+}
+            `,
+    );
+
+    await db.$unuseAll().user.create({
+        data: { id: 1, secret: 'secret', posts: { create: { id: 1, published: true } } },
+    });
+    await db.$unuseAll().user.create({
+        data: { id: 2, secret: 'secret' },
+    });
+
+    await expect(db.user.findFirst({ where: { id: 1 } })).resolves.toMatchObject({ secret: 'secret' });
+    await expect(db.user.findFirst({ where: { id: 1 }, select: { id: true } })).resolves.toEqual({ id: 1 });
+
+    let r = await db.user.findFirst({ where: { id: 2 } });
+    expect(r.secret).toBeUndefined();
+    r = await db.user.findFirst({ where: { id: 2 }, select: { id: true } });
+    expect(r.secret).toBeUndefined();
+});

tests/regression/test/v2-migrated/issue-1984.test.ts

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