fix(policy): update fails for model using both @password and @@validate

fixes #2000

Author: ymc9

packages/runtime/src/enhancements/node/policy/policy-utils.ts

@@ -826,6 +826,8 @@ export class PolicyUtil extends QueryUtils {
     /**
      * Given a model and a unique filter, checks the operation is allowed by policies and field validations.
      * Rejects with an error if not allowed.
+     *
+     * This method is only called by mutation operations.
      */
     async checkPolicyForUnique(
         model: string,
@@ -1365,32 +1367,68 @@ export class PolicyUtil extends QueryUtils {
         excludePasswordFields: boolean = true,
         kind: 'create' | 'update' | undefined = undefined
     ) {
+        if (!this.zodSchemas) {
+            return undefined;
+        }
+
         if (!this.hasFieldValidation(model)) {
             return undefined;
         }
+
         const schemaKey = `${upperCaseFirst(model)}${kind ? 'Prisma' + upperCaseFirst(kind) : ''}Schema`;
-        let result = this.zodSchemas?.models?.[schemaKey] as ZodObject<any> | undefined;
-
-        if (result && excludePasswordFields) {
-            // fields with `@password` attribute changes at runtime, so we cannot directly use the generated
-            // zod schema to validate it, instead, the validation happens when checking the input of "create"
-            // and "update" operations
-            const modelFields = this.modelMeta.models[lowerCaseFirst(model)]?.fields;
-            if (modelFields) {
-                for (const [key, field] of Object.entries(modelFields)) {
-                    if (field.attributes?.some((attr) => attr.name === '@password')) {
-                        // override `@password` field schema with a string schema
-                        let pwFieldSchema: ZodSchema = z.string();
-                        if (field.isOptional) {
-                            pwFieldSchema = pwFieldSchema.nullish();
+
+        if (excludePasswordFields) {
+            // The `excludePasswordFields` mode is to handle the issue the fields marked with `@password` change at runtime,
+            // so they can only be fully validated when processing the input of "create" and "update" operations.
+            //
+            // When excluding them, we need to override them with plain string schemas. However, since the scheme is not always
+            // an `ZodObject` (this happens when there's `@@validate` refinement), we need to fetch the `ZodObject` schema before
+            // the refinement is applied, override the `@password` fields and then re-apply the refinement.
+
+            let schema: ZodObject<any> | undefined;
+
+            const overridePasswordFields = (schema: z.ZodObject<any>) => {
+                let result = schema;
+                const modelFields = this.modelMeta.models[lowerCaseFirst(model)]?.fields;
+                if (modelFields) {
+                    for (const [key, field] of Object.entries(modelFields)) {
+                        if (field.attributes?.some((attr) => attr.name === '@password')) {
+                            // override `@password` field schema with a string schema
+                            let pwFieldSchema: ZodSchema = z.string();
+                            if (field.isOptional) {
+                                pwFieldSchema = pwFieldSchema.nullish();
+                            }
+                            result = result.merge(z.object({ [key]: pwFieldSchema }));
                         }
-                        result = result?.merge(z.object({ [key]: pwFieldSchema }));
                     }
                 }
+                return result;
+            };
+
+            // get the schema without refinement: `[Model]WithoutRefineSchema`
+            const withoutRefineSchemaKey = `${upperCaseFirst(model)}${
+                kind ? 'Prisma' + upperCaseFirst(kind) : ''
+            }WithoutRefineSchema`;
+            schema = this.zodSchemas.models[withoutRefineSchemaKey] as ZodObject<any> | undefined;
+
+            if (schema) {
+                // the schema has refinement, need to call refine function after schema merge
+                schema = overridePasswordFields(schema);
+                // refine function: `refine[Model]`
+                const refineFuncKey = `refine${upperCaseFirst(model)}`;
+                const refineFunc = this.zodSchemas.models[refineFuncKey] as unknown as (
+                    schema: ZodObject<any>
+                ) => ZodSchema;
+                return typeof refineFunc === 'function' ? refineFunc(schema) : schema;
+            } else {
+                // otherwise, directly override the `@password` fields
+                schema = this.zodSchemas.models[schemaKey] as ZodObject<any> | undefined;
+                return schema ? overridePasswordFields(schema) : undefined;
             }
+        } else {
+            // simply return the schema
+            return this.zodSchemas.models[schemaKey];
         }
-
-        return result;
     }
 
     /**

tests/regression/tests/issue-2000.test.ts

@@ -0,0 +1,69 @@
+import { loadSchema } from '@zenstackhq/testtools';
+
+describe('issue 2000', () => {
+    it('regression', async () => {
+        const { enhance } = await loadSchema(
+            `
+            abstract model Base {
+                id              String          @id  @default(uuid())    @deny('update', true)
+                createdAt       DateTime        @default(now())          @deny('update', true)
+                updatedAt       DateTime        @updatedAt               @deny('update', true)
+                active          Boolean         @default(false)
+                published       Boolean         @default(true)           
+                deleted         Boolean         @default(false)
+                startDate       DateTime?
+                endDate         DateTime?
+
+                @@allow('create', true)
+                @@allow('read', true)
+                @@allow('update', true)
+            }
+
+            enum EntityType {
+                User
+                Alias
+                Group
+                Service
+                Device
+                Organization
+                Guest
+            }
+
+            model Entity extends Base {
+                entityType      EntityType
+                name            String?                 @unique
+                members         Entity[]                @relation("members")
+                memberOf        Entity[]                @relation("members")
+                @@delegate(entityType)
+
+
+                @@allow('create', true)
+                @@allow('read', true)
+                @@allow('update', true)
+                @@validate(!active || (active && name != null), "Active Entities Must Have A Name")
+            }
+
+            model User extends Entity {
+                profile         Json?               
+                username        String                  @unique 
+                password        String                  @password
+
+                @@allow('create', true)
+                @@allow('read', true)
+                @@allow('update', true)
+            }
+            `
+        );
+
+        const db = enhance();
+        await expect(db.user.create({ data: { username: 'admin', password: 'abc12345' } })).toResolveTruthy();
+        await expect(
+            db.user.update({ where: { username: 'admin' }, data: { password: 'abc123456789123' } })
+        ).toResolveTruthy();
+
+        // violating validation rules
+        await expect(
+            await db.user.update({ where: { username: 'admin' }, data: { active: true } })
+        ).toBeRejectedByPolicy();
+    });
+});