feat: add @@auth option for declaring auth model (zenstackhq#787)

Co-authored-by: Jason MacDonald <jason.macdonald@veeva.com> Co-authored-by: ymc9 <104139426+ymc9@users.noreply.github.com>
Eventiva · Nov 1, 2023 · c390de1 · c390de1
1 parent 42d654f
commit c390de1
Show file tree

Hide file tree

Showing 5 changed files with 428 additions and 115 deletions.
diff --git a/packages/schema/src/language-server/zmodel-linker.ts b/packages/schema/src/language-server/zmodel-linker.ts
@@ -36,7 +36,7 @@ import {
     isReferenceExpr,
     isStringLiteral,
 } from '@zenstackhq/language/ast';
-import { getContainingModel, isFromStdlib } from '@zenstackhq/sdk';
+import { getContainingModel, hasAttribute, isFromStdlib } from '@zenstackhq/sdk';
 import {
     AstNode,
     AstNodeDescription,
@@ -278,9 +278,15 @@ export class ZModelLinker extends DefaultLinker {
                 const model = getContainingModel(node);
 
                 if (model) {
-                    const userModel = getAllDeclarationsFromImports(this.langiumDocuments(), model).find(
-                        (d) => isDataModel(d) && d.name === 'User'
-                    );
+                    let userModel;
+                    userModel = getAllDeclarationsFromImports(this.langiumDocuments(), model).find((d) => {
+                        return isDataModel(d) && hasAttribute(d, '@@auth');
+                    });
+                    if (!userModel) {
+                        userModel = getAllDeclarationsFromImports(this.langiumDocuments(), model).find((d) => {
+                            return isDataModel(d) && d.name === 'User';
+                        });
+                    }
                     if (userModel) {
                         node.$resolvedType = { decl: userModel, nullable: true };
                     }

diff --git a/packages/schema/src/res/stdlib.zmodel b/packages/schema/src/res/stdlib.zmodel
@@ -376,6 +376,11 @@ attribute @@deny(_ operation: String, _ condition: Boolean)
  */
 attribute @deny(_ operation: String, _ condition: Boolean)
 
+/**
+ * Defines the model to use when when checking access policies.
+ */
+attribute @@auth()
+
 /**
  * Indicates that the field is a password field and needs to be hashed before persistence.
  * 

diff --git a/packages/schema/tests/generator/expression-writer.test.ts b/packages/schema/tests/generator/expression-writer.test.ts
@@ -883,6 +883,40 @@ describe('Expression Writer Tests', () => {
         );
     });
 
+    it('auth using different model check', async () => {
+        await check(
+            `
+            model Membership {
+                id String @id
+                t Test?
+                @@auth()
+            }
+
+            model Test {
+                id String @id
+                owner Membership @relation(fields: [ownerId], references: [id])
+                ownerId String @unique @allow('all', auth().id == owner.id)
+                value Int
+                @@allow('all', auth().id == owner.id)
+               
+            }
+                `,
+            (model) => {
+                const args = model.attributes[0].args[1];
+                return args.value;
+            },
+            `((user?.id??null)==null)?{
+                OR:[]
+            }:{
+                owner:{
+                     id:{
+                         equals:(user?.id??null)
+                    }
+                }
+            }`
+        );
+    });
+
     it('relation field null check', async () => {
         await check(
             `
@@ -903,8 +937,7 @@ describe('Expression Writer Tests', () => {
             `
             {
                 OR: [{ m: { is: null } }, { m: { s: { equals: null } } }]
-            }
-            `
+            }`
         );
 
         await check(

diff --git a/packages/sdk/src/utils.ts b/packages/sdk/src/utils.ts
@@ -111,7 +111,7 @@ export function hasAttribute(decl: DataModel | DataModelField | Enum | EnumField
 
 export function getAttribute(decl: DataModel | DataModelField | Enum | EnumField, name: string) {
     return (decl.attributes as (DataModelAttribute | DataModelFieldAttribute)[]).find(
-        (attr) => resolved(attr.decl).name === name
+        (attr) => attr.decl.$refText === name
     );
 }