feat: support using collection predicate expression with auth() (#831)

Author: ymc9

packages/runtime/src/enhancements/policy/index.ts

@@ -10,6 +10,7 @@ import { hasAllFields } from '../../validation';
 import { makeProxy } from '../proxy';
 import type { CommonEnhancementOptions, PolicyDef, ZodSchemas } from '../types';
 import { PolicyProxyHandler } from './handler';
+import { Logger } from './logger';
 
 /**
  * Context for evaluating access policies
@@ -72,7 +73,8 @@ export function withPolicy<DbClient extends object>(
     const _zodSchemas = options?.zodSchemas ?? getDefaultZodSchemas(options?.loadPath);
 
     // validate user context
-    if (context?.user && _modelMeta.authModel) {
+    const userContext = context?.user;
+    if (userContext && _modelMeta.authModel) {
         const idFields = getIdFields(_modelMeta, _modelMeta.authModel);
         if (
             !hasAllFields(
@@ -84,6 +86,16 @@ export function withPolicy<DbClient extends object>(
                 `Invalid user context: must have valid ID field ${idFields.map((f) => `"${f.name}"`).join(', ')}`
             );
         }
+
+        // validate user context for fields used in policy expressions
+        const authSelector = _policy.authSelector;
+        if (authSelector) {
+            Object.keys(authSelector).forEach((f) => {
+                if (!(f in userContext)) {
+                    console.warn(`User context does not have field "${f}" used in policy rules`);
+                }
+            });
+        }
     }
 
     return makeProxy(

packages/runtime/src/enhancements/types.ts

@@ -56,7 +56,12 @@ export type PolicyDef = {
                 [HAS_FIELD_LEVEL_POLICY_FLAG]?: boolean;
             }
     >;
+
+    // tracks which models have data validation rules
     validation: Record<string, { hasValidation: boolean }>;
+
+    // a { select: ... } object for fetching `auth()` fields needed for policy evaluation
+    authSelector?: object;
 };
 
 /**

packages/schema/src/plugins/access-policy/expression-writer.ts

@@ -226,8 +226,12 @@ export class ExpressionWriter {
         // check if the operand should be compiled to a relation query
         // or a plain expression
         const compileToRelationQuery =
-            (this.isPostGuard && this.isFutureMemberAccess(expr.left)) ||
-            (!this.isPostGuard && !this.isFutureMemberAccess(expr.left));
+            // expression rooted to `auth()` is always compiled to plain expression
+            !this.isAuthOrAuthMemberAccess(expr.left) &&
+            // `future()` in post-update context
+            ((this.isPostGuard && this.isFutureMemberAccess(expr.left)) ||
+                // non-`future()` in pre-update context
+                (!this.isPostGuard && !this.isFutureMemberAccess(expr.left)));
 
         if (compileToRelationQuery) {
             this.block(() => {

packages/schema/src/plugins/access-policy/policy-guard-generator.ts

@@ -81,7 +81,6 @@ export default class PolicyGenerator {
             namedImports: [
                 { name: 'type QueryContext' },
                 { name: 'type DbOperations' },
-                { name: 'hasAllFields' },
                 { name: 'allFieldsEqual' },
                 { name: 'type PolicyDef' },
             ],
@@ -104,6 +103,8 @@ export default class PolicyGenerator {
             policyMap[model.name] = await this.generateQueryGuardForModel(model, sf);
         }
 
+        const authSelector = this.generateAuthSelector(models);
+
         sf.addVariableStatement({
             declarationKind: VariableDeclarationKind.Const,
             declarations: [
@@ -140,6 +141,11 @@ export default class PolicyGenerator {
                                     writer.writeLine(',');
                                 }
                             });
+
+                            if (authSelector) {
+                                writer.writeLine(',');
+                                writer.write(`authSelector: ${JSON.stringify(authSelector)}`);
+                            }
                         });
                     },
                 },
@@ -165,6 +171,43 @@ export default class PolicyGenerator {
         }
     }
 
+    // Generates a { select: ... } object to select `auth()` fields used in policy rules
+    private generateAuthSelector(models: DataModel[]) {
+        const authRules: Expression[] = [];
+
+        models.forEach((model) => {
+            // model-level rules
+            const modelPolicyAttrs = model.attributes.filter((attr) =>
+                ['@@allow', '@@deny'].includes(attr.decl.$refText)
+            );
+
+            // field-level rules
+            const fieldPolicyAttrs = model.fields
+                .flatMap((f) => f.attributes)
+                .filter((attr) => ['@allow', '@deny'].includes(attr.decl.$refText));
+
+            // all rule expression
+            const allExpressions = [...modelPolicyAttrs, ...fieldPolicyAttrs]
+                .filter((attr) => attr.args.length > 1)
+                .map((attr) => attr.args[1].value);
+
+            // collect `auth()` member access
+            allExpressions.forEach((rule) => {
+                streamAst(rule).forEach((node) => {
+                    if (isMemberAccessExpr(node) && isAuthInvocation(node.operand)) {
+                        authRules.push(node);
+                    }
+                });
+            });
+        });
+
+        if (authRules.length > 0) {
+            return this.generateSelectForRules(authRules, true);
+        } else {
+            return undefined;
+        }
+    }
+
     private isEnumReferenced(model: Model, decl: Enum): unknown {
         return streamAllContents(model).some((node) => {
             if (isDataModelField(node) && node.type.reference?.ref === decl) {
@@ -293,7 +336,7 @@ export default class PolicyGenerator {
             result[kind] = guardFunc.getName()!;
 
             if (kind === 'postUpdate') {
-                const preValueSelect = this.generateSelectForRules(allows, denies);
+                const preValueSelect = this.generateSelectForRules([...allows, ...denies]);
                 if (preValueSelect) {
                     result[PRE_UPDATE_VALUE_SELECTOR] = preValueSelect;
                 }
@@ -340,7 +383,7 @@ export default class PolicyGenerator {
 
         if (allFieldsAllows.length > 0 || allFieldsDenies.length > 0) {
             result[HAS_FIELD_LEVEL_POLICY_FLAG] = true;
-            const readFieldCheckSelect = this.generateSelectForRules(allFieldsAllows, allFieldsDenies);
+            const readFieldCheckSelect = this.generateSelectForRules([...allFieldsAllows, ...allFieldsDenies]);
             if (readFieldCheckSelect) {
                 result[FIELD_LEVEL_READ_CHECKER_SELECTOR] = readFieldCheckSelect;
             }
@@ -477,7 +520,7 @@ export default class PolicyGenerator {
 
     // generates a "select" object that contains (recursively) fields referenced by the
     // given policy rules
-    private generateSelectForRules(allows: Expression[], denies: Expression[]): object {
+    private generateSelectForRules(rules: Expression[], forAuthContext = false): object {
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
         const result: any = {};
         const addPath = (path: string[]) => {
@@ -504,6 +547,10 @@ export default class PolicyGenerator {
                     return [target.name];
                 }
             } else if (isMemberAccessExpr(node)) {
+                if (forAuthContext && isAuthInvocation(node.operand)) {
+                    return [node.member.$refText];
+                }
+
                 if (isFutureExpr(node.operand)) {
                     // future().field is not subject to pre-update select
                     return undefined;
@@ -562,7 +609,7 @@ export default class PolicyGenerator {
             }
         };
 
-        for (const rule of [...allows, ...denies]) {
+        for (const rule of rules) {
             const paths = collectReferencePaths(rule);
             paths.forEach((p) => addPath(p));
         }
@@ -780,11 +827,7 @@ export default class PolicyGenerator {
             }
 
             // normalize user to null to avoid accidentally use undefined in filter
-            statements.push(
-                `const user = hasAllFields(context.user, [${userIdFields
-                    .map((f) => "'" + f.name + "'")
-                    .join(', ')}]) ? context.user as any : null;`
-            );
+            statements.push(`const user: any = context.user ?? null;`);
         }
     }
 }

packages/schema/src/utils/ast-utils.ts

@@ -153,6 +153,6 @@ export function getAllDeclarationsFromImports(documents: LangiumDocuments, model
     return model.declarations.concat(...imports.map((imp) => imp.declarations));
 }
 
-export function isCollectionPredicate(expr: Expression): expr is BinaryExpr {
-    return isBinaryExpr(expr) && ['?', '!', '^'].includes(expr.operator);
+export function isCollectionPredicate(node: AstNode): node is BinaryExpr {
+    return isBinaryExpr(node) && ['?', '!', '^'].includes(node.operator);
 }

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

@@ -212,4 +212,100 @@ describe('With Policy: auth() test', () => {
         const adminDb = withPolicy({ id: 'user1', role: 'ADMIN' });
         await expect(adminDb.post.create({ data: { title: 'abc' } })).toResolveTruthy();
     });
+
+    it('collection predicate', async () => {
+        const { enhance, prisma } = await loadSchema(
+            `
+        model User {
+            id String @id @default(uuid())
+            posts Post[]
+
+            @@allow('all', true)
+        }
+
+        model Post {
+            id String @id @default(uuid())
+            title String
+            published Boolean @default(false)
+            author User @relation(fields: [authorId], references: [id])
+            authorId String
+            comments Comment[]
+
+            @@allow('read', true)
+            @@allow('create', auth().posts?[published && comments![published]])
+        }
+
+        model Comment {
+            id String @id @default(uuid())
+            published Boolean @default(false)
+            post Post @relation(fields: [postId], references: [id])
+            postId String
+
+            @@allow('all', true)
+        }
+        `
+        );
+
+        const user = await prisma.user.create({ data: {} });
+
+        const createPayload = {
+            data: { title: 'Post 1', author: { connect: { id: user.id } } },
+        };
+
+        // no post
+        await expect(enhance({ id: '1' }).post.create(createPayload)).toBeRejectedByPolicy();
+
+        // post not published
+        await expect(
+            enhance({ id: '1', posts: [{ id: '1', published: false }] }).post.create(createPayload)
+        ).toBeRejectedByPolicy();
+
+        // no comments
+        await expect(
+            enhance({ id: '1', posts: [{ id: '1', published: true }] }).post.create(createPayload)
+        ).toBeRejectedByPolicy();
+
+        // not all comments published
+        await expect(
+            enhance({
+                id: '1',
+                posts: [
+                    {
+                        id: '1',
+                        published: true,
+                        comments: [
+                            { id: '1', published: true },
+                            { id: '2', published: false },
+                        ],
+                    },
+                ],
+            }).post.create(createPayload)
+        ).toBeRejectedByPolicy();
+
+        // comments published but parent post is not
+        await expect(
+            enhance({
+                id: '1',
+                posts: [
+                    { id: '1', published: false, comments: [{ id: '1', published: true }] },
+                    { id: '2', published: true },
+                ],
+            }).post.create(createPayload)
+        ).toBeRejectedByPolicy();
+
+        await expect(
+            enhance({
+                id: '1',
+                posts: [
+                    { id: '1', published: true, comments: [{ id: '1', published: true }] },
+                    { id: '2', published: false },
+                ],
+            }).post.create(createPayload)
+        ).toResolveTruthy();
+
+        // no comments ("every" evaluates to tru in this case)
+        await expect(
+            enhance({ id: '1', posts: [{ id: '1', published: true, comments: [] }] }).post.create(createPayload)
+        ).toResolveTruthy();
+    });
 });