From 6e9d24254cee5e927c611d40224738d6338a73cc Mon Sep 17 00:00:00 2001
From: Yiming <yiming@whimslab.io>
Date: Mon, 25 Mar 2024 15:10:02 -0700
Subject: [PATCH] chore: disallow `this` in collection predicate (#1176)

---
 .../validator/expression-validator.ts         | 29 ++++++++++++++++---
 .../validation/datamodel-validation.test.ts   | 23 +++++++++++++++
 .../tests/regression/issue-925.test.ts        |  3 +-
 .../tests/regression/issues.test.ts           |  2 +-
 4 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/packages/schema/src/language-server/validator/expression-validator.ts b/packages/schema/src/language-server/validator/expression-validator.ts
index 8a87ddc14..59ad7edbb 100644
--- a/packages/schema/src/language-server/validator/expression-validator.ts
+++ b/packages/schema/src/language-server/validator/expression-validator.ts
@@ -13,8 +13,8 @@ import {
     isThisExpr,
 } from '@zenstackhq/language/ast';
 import { isAuthInvocation, isDataModelFieldReference, isEnumFieldReference } from '@zenstackhq/sdk';
-import { ValidationAcceptor } from 'langium';
-import { findUpAst, getContainingDataModel, isCollectionPredicate } from '../../utils/ast-utils';
+import { ValidationAcceptor, streamAst } from 'langium';
+import { findUpAst, getContainingDataModel } from '../../utils/ast-utils';
 import { AstValidator } from '../types';
 import { typeAssignable } from './utils';
 
@@ -32,8 +32,6 @@ export default class ExpressionValidator implements AstValidator<Expression> {
                     'auth() cannot be resolved because no model marked wth "@@auth()" or named "User" is found',
                     { node: expr }
                 );
-            } else if (isCollectionPredicate(expr)) {
-                accept('error', 'collection predicate can only be used on an array of model type', { node: expr });
             } else {
                 accept('error', 'expression cannot be resolved', {
                     node: expr,
@@ -221,6 +219,29 @@ export default class ExpressionValidator implements AstValidator<Expression> {
                 }
                 break;
             }
+
+            case '?':
+            case '!':
+            case '^':
+                this.validateCollectionPredicate(expr, accept);
+                break;
+        }
+    }
+
+    private validateCollectionPredicate(expr: BinaryExpr, accept: ValidationAcceptor) {
+        if (!expr.$resolvedType) {
+            accept('error', 'collection predicate can only be used on an array of model type', { node: expr });
+            return;
+        }
+
+        // TODO: revisit this when we implement lambda inside collection predicate
+        const thisExpr = streamAst(expr).find(isThisExpr);
+        if (thisExpr) {
+            accept(
+                'error',
+                'using `this` in collection predicate is not supported. To compare entity identity, use id field comparison instead.',
+                { node: thisExpr }
+            );
         }
     }
 
diff --git a/packages/schema/tests/schema/validation/datamodel-validation.test.ts b/packages/schema/tests/schema/validation/datamodel-validation.test.ts
index 4f44b109e..0bf12245d 100644
--- a/packages/schema/tests/schema/validation/datamodel-validation.test.ts
+++ b/packages/schema/tests/schema/validation/datamodel-validation.test.ts
@@ -78,6 +78,29 @@ describe('Data Model Validation Tests', () => {
         ).toMatchObject(errorLike('Field of "Unsupported" type cannot be used in expressions'));
     });
 
+    it('Using `this` in collection predicate', async () => {
+        expect(
+            await safelyLoadModel(`
+            ${prelude}
+            model User {
+                id String @id
+                members User[]
+                @@allow('all', members?[this == auth()])
+            }
+        `)
+        ).toMatchObject(errorLike('using `this` in collection predicate is not supported'));
+
+        expect(
+            await loadModel(`
+        model User {
+            id String @id
+            members User[]
+            @@allow('all', members?[id == auth().id])
+        }
+        `)
+        ).toBeTruthy();
+    });
+
     it('mix array and optional', async () => {
         expect(
             await safelyLoadModel(`
diff --git a/tests/integration/tests/regression/issue-925.test.ts b/tests/integration/tests/regression/issue-925.test.ts
index b19d9d615..19ef210bf 100644
--- a/tests/integration/tests/regression/issue-925.test.ts
+++ b/tests/integration/tests/regression/issue-925.test.ts
@@ -35,7 +35,8 @@ describe('Regression: issue 925', () => {
         ).resolves.toContain("Could not resolve reference to ReferenceTarget named 'test'.");
     });
 
-    it('reference with this', async () => {
+    // eslint-disable-next-line jest/no-disabled-tests
+    it.skip('reference with this', async () => {
         await loadModel(
             `
             model User {
diff --git a/tests/integration/tests/regression/issues.test.ts b/tests/integration/tests/regression/issues.test.ts
index 7c2ca94cd..318682aad 100644
--- a/tests/integration/tests/regression/issues.test.ts
+++ b/tests/integration/tests/regression/issues.test.ts
@@ -600,7 +600,7 @@ model User {
     // can be created by anyone, even not logged in
     @@allow('create', true)
     // can be read by users in the same organization
-    @@allow('read', orgs?[members?[auth() == this]])
+    @@allow('read', orgs?[members?[auth().id == id]])
     // full access by oneself
     @@allow('all', auth() == this)
 }