From a3596188f79aab8d84dfc1ce35fce313e8b29f97 Mon Sep 17 00:00:00 2001
From: ymc9 <104139426+ymc9@users.noreply.github.com>
Date: Mon, 21 Aug 2023 22:42:14 +0800
Subject: [PATCH] fix: typing of policy definition

---
 .../src/enhancements/policy/policy-utils.ts   |  4 ++--
 packages/runtime/src/enhancements/types.ts    | 23 +++++++++++--------
 .../access-policy/policy-guard-generator.ts   | 11 +++++++--
 3 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/packages/runtime/src/enhancements/policy/policy-utils.ts b/packages/runtime/src/enhancements/policy/policy-utils.ts
index 142ab154a..cf77fe4b0 100644
--- a/packages/runtime/src/enhancements/policy/policy-utils.ts
+++ b/packages/runtime/src/enhancements/policy/policy-utils.ts
@@ -18,7 +18,7 @@ import { AuthUser, DbClientContract, DbOperations, FieldInfo, PolicyOperationKin
 import { getVersion } from '../../version';
 import { getFields, resolveField } from '../model-meta';
 import { NestedWriteVisitorContext } from '../nested-write-vistor';
-import type { InputCheckFunc, ModelMeta, PolicyDef, PolicyFunc, ReadFieldCheckFunc, ZodSchemas } from '../types';
+import type { InputCheckFunc, ModelMeta, PolicyDef, ReadFieldCheckFunc, ZodSchemas } from '../types';
 import {
     formatObject,
     getIdFields,
@@ -223,7 +223,7 @@ export class PolicyUtil {
         if (!guard) {
             return false;
         }
-        const provider: PolicyFunc | boolean | undefined = guard[operation];
+        const provider = guard[operation];
         return typeof provider !== 'boolean' || provider !== true;
     }
 
diff --git a/packages/runtime/src/enhancements/types.ts b/packages/runtime/src/enhancements/types.ts
index d879bf510..a3d0e6a6f 100644
--- a/packages/runtime/src/enhancements/types.ts
+++ b/packages/runtime/src/enhancements/types.ts
@@ -1,11 +1,13 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import { z } from 'zod';
+import type { DbOperations, FieldInfo, PolicyOperationKind, QueryContext } from '../types';
 import {
     FIELD_LEVEL_READ_CHECKER_SELECTOR,
-    HAS_FIELD_LEVEL_POLICY_FLAG,
     PRE_UPDATE_VALUE_SELECTOR,
+    FIELD_LEVEL_READ_CHECKER_PREFIX,
+    FIELD_LEVEL_UPDATE_GUARD_PREFIX,
+    HAS_FIELD_LEVEL_POLICY_FLAG,
 } from '../constants';
-import type { DbOperations, FieldInfo, PolicyOperationKind, QueryContext } from '../types';
 
 /**
  * Metadata for a model-level unique constraint
@@ -43,15 +45,18 @@ export type PolicyDef = {
     // Prisma query guards
     guard: Record<
         string,
-        {
-            allowAll?: boolean;
-            denyAll?: boolean;
-        } & Partial<Record<PolicyOperationKind, PolicyFunc>> & {
-                create_input: InputCheckFunc;
-            } & {
+        // policy operation guard functions
+        Partial<Record<PolicyOperationKind, PolicyFunc | boolean>> &
+            // 'create_input' checker function
+            Partial<Record<`${PolicyOperationKind}_input`, InputCheckFunc | boolean>> &
+            // field-level read checker functions or update guard functions
+            Record<`${typeof FIELD_LEVEL_READ_CHECKER_PREFIX}${string}`, ReadFieldCheckFunc> &
+            Record<`${typeof FIELD_LEVEL_UPDATE_GUARD_PREFIX}${string}`, PolicyFunc> & {
+                // pre-update value selector
                 [PRE_UPDATE_VALUE_SELECTOR]?: object;
+                // field-level read checker selector
                 [FIELD_LEVEL_READ_CHECKER_SELECTOR]?: object;
-            } & Record<string, ReadFieldCheckFunc | PolicyFunc> & {
+                // flag that indicates if the model has field-level access control
                 [HAS_FIELD_LEVEL_POLICY_FLAG]?: boolean;
             }
     >;
diff --git a/packages/schema/src/plugins/access-policy/policy-guard-generator.ts b/packages/schema/src/plugins/access-policy/policy-guard-generator.ts
index 28ac12805..870dc40df 100644
--- a/packages/schema/src/plugins/access-policy/policy-guard-generator.ts
+++ b/packages/schema/src/plugins/access-policy/policy-guard-generator.ts
@@ -22,6 +22,7 @@ import {
     FIELD_LEVEL_READ_CHECKER_SELECTOR,
     FIELD_LEVEL_UPDATE_GUARD_PREFIX,
     HAS_FIELD_LEVEL_POLICY_FLAG,
+    PRE_UPDATE_VALUE_SELECTOR,
     type PolicyKind,
     type PolicyOperationKind,
 } from '@zenstackhq/runtime';
@@ -74,7 +75,12 @@ export default class PolicyGenerator {
         sf.addStatements('/* eslint-disable */');
 
         sf.addImportDeclaration({
-            namedImports: [{ name: 'type QueryContext' }, { name: 'type DbOperations' }, { name: 'hasAllFields' }],
+            namedImports: [
+                { name: 'type QueryContext' },
+                { name: 'type DbOperations' },
+                { name: 'hasAllFields' },
+                { name: 'type PolicyDef' },
+            ],
             moduleSpecifier: `${RUNTIME_PACKAGE}`,
         });
 
@@ -99,6 +105,7 @@ export default class PolicyGenerator {
             declarations: [
                 {
                     name: 'policy',
+                    type: 'PolicyDef',
                     initializer: (writer) => {
                         writer.block(() => {
                             writer.write('guard:');
@@ -256,7 +263,7 @@ export default class PolicyGenerator {
             if (kind === 'postUpdate') {
                 const preValueSelect = this.generateSelectForRules(allows, denies);
                 if (preValueSelect) {
-                    result['preValueSelect'] = preValueSelect;
+                    result[PRE_UPDATE_VALUE_SELECTOR] = preValueSelect;
                 }
             }