function-invocation-validator.ts

import {
    AbstractCallable,
    AliasDecl,
    Argument,
    DataModel,
    DataModelAttribute,
    DataModelFieldAttribute,
    Expression,
    FunctionParam,
    InvocationExpr,
    isAliasDecl,
    isArrayExpr,
    isDataModel,
    isDataModelAttribute,
    isDataModelFieldAttribute,
    isInvocationExpr,
    isLiteralExpr,
} from '@zenstackhq/language/ast';
import {
    ExpressionContext,
    getFieldReference,
    getFunctionExpressionContext,
    getLiteral,
    isDataModelFieldReference,
    isEnumFieldReference,
    isFromStdlib,
    isValidationAttribute,
} from '@zenstackhq/sdk';
import { AstNode, streamAst, ValidationAcceptor } from 'langium';
import { match, P } from 'ts-pattern';
import { isCheckInvocation } from '../../utils/ast-utils';
import { AstValidator } from '../types';
import { isAuthOrAuthMemberAccess, typeAssignable } from './utils';

// a registry of function handlers marked with @func
const invocationCheckers = new Map<string, PropertyDescriptor>();

// function handler decorator
function func(name: string) {
    return function (_target: unknown, _propertyKey: string, descriptor: PropertyDescriptor) {
        if (!invocationCheckers.get(name)) {
            invocationCheckers.set(name, descriptor);
        }
        return descriptor;
    };
}
/**
 * InvocationExpr validation
 */
export default class FunctionInvocationValidator implements AstValidator<Expression> {
    validate(expr: InvocationExpr, accept: ValidationAcceptor): void {
        const callableDecl = expr.function.ref;
        if (!callableDecl) {
            accept('error', 'function or alias cannot be resolved', { node: expr });
            return;
        }

        if (!this.validateArgs(callableDecl, expr.args, accept)) {
            return;
        }

        if (isFromStdlib(callableDecl)) {
            // validate standard library functions

            // find the containing attribute context for the invocation
            let curr: AstNode | undefined = expr.$container;
            let containerAttribute: DataModelAttribute | DataModelFieldAttribute | AliasDecl | undefined;
            while (curr) {
                if (isDataModelAttribute(curr) || isDataModelFieldAttribute(curr) || isAliasDecl(curr)) {
                    containerAttribute = curr;
                    break;
                }
                curr = curr.$container;
            }

            // validate the context allowed for the function
            const exprContext = this.getExpressionContext(containerAttribute);

            // get the context allowed for the function
            const funcAllowedContext = getFunctionExpressionContext(callableDecl);

            if (funcAllowedContext.length > 0 && (!exprContext || !funcAllowedContext.includes(exprContext))) {
                accept(
                    'error',
                    `function "${callableDecl.name}" is not allowed in the current context${
                        exprContext ? ': ' + exprContext : ''
                    }`,
                    {
                        node: expr,
                    }
                );
                return;
            }

            // TODO: express function validation rules declaratively in ZModel

            const allCasing = ['original', 'upper', 'lower', 'capitalize', 'uncapitalize'];
            if (['currentModel', 'currentOperation'].includes(callableDecl.name)) {
                const arg = getLiteral<string>(expr.args[0]?.value);
                if (arg && !allCasing.includes(arg)) {
                    accept('error', `argument must be one of: ${allCasing.map((c) => '"' + c + '"').join(', ')}`, {
                        node: expr.args[0],
                    });
                }
            } else if (
                funcAllowedContext.includes(ExpressionContext.AccessPolicy) ||
                funcAllowedContext.includes(ExpressionContext.ValidationRule)
            ) {
                // filter operation functions validation

                // first argument must refer to a model field
                const firstArg = expr.args?.[0]?.value;
                if (firstArg) {
                    if (!getFieldReference(firstArg)) {
                        accept('error', 'first argument must be a field reference', { node: firstArg });
                    }
                }

                // second argument must be a literal or array of literal
                const secondArg = expr.args?.[1]?.value;
                if (
                    secondArg &&
                    // literal
                    !isLiteralExpr(secondArg) &&
                    // enum field
                    !isEnumFieldReference(secondArg) &&
                    // `auth()...` expression
                    !isAuthOrAuthMemberAccess(secondArg) &&
                    // static function calls that are runtime constants: `currentModel`, `currentOperation`
                    !this.isStaticFunctionCall(secondArg) &&
                    // array of literal/enum
                    !(
                        isArrayExpr(secondArg) &&
                        secondArg.items.every(
                            (item: Expression) =>
                                isLiteralExpr(item) || isEnumFieldReference(item) || isAuthOrAuthMemberAccess(item)
                        )
                    )
                ) {
                    accept(
                        'error',
                        'second argument must be a literal, an enum, an expression starting with `auth().`, or an array of them',
                        {
                            node: secondArg,
                        }
                    );
                }
            }

            // run checkers for specific functions
            const checker = invocationCheckers.get(expr.function.$refText);
            if (checker) {
                checker.value.call(this, expr, accept);
            }
        }
    }

    private getExpressionContext(
        containerAttribute: DataModelAttribute | DataModelFieldAttribute | AliasDecl | undefined
    ) {
        if (!containerAttribute) {
            return undefined;
        }
        if (isAliasDecl(containerAttribute)) {
            return ExpressionContext.AliasFunction;
        }
        if (isValidationAttribute(containerAttribute)) {
            return ExpressionContext.ValidationRule;
        }
        return match(containerAttribute?.decl.$refText)
            .with('@default', () => ExpressionContext.DefaultValue)
            .with(P.union('@@allow', '@@deny', '@allow', '@deny'), () => ExpressionContext.AccessPolicy)
            .with('@@index', () => ExpressionContext.Index)
            .otherwise(() => undefined);
    }

    private isStaticFunctionCall(expr: Expression) {
        return isInvocationExpr(expr) && ['currentModel', 'currentOperation'].includes(expr.function.$refText);
    }

    private validateArgs(funcDecl: AbstractCallable, args: Argument[], accept: ValidationAcceptor) {
        let success = true;
        for (let i = 0; i < funcDecl.params.length; i++) {
            const param = funcDecl.params[i];
            const arg = args[i];
            if (!arg) {
                if (!param.optional) {
                    accept('error', `missing argument for parameter "${param.name}"`, { node: funcDecl });
                    success = false;
                }
            } else {
                if (!this.validateInvocationArg(arg, param, accept)) {
                    success = false;
                }
            }
        }
        // TODO: do we need to complain for extra arguments?
        return success;
    }

    private validateInvocationArg(arg: Argument, param: FunctionParam, accept: ValidationAcceptor) {
        const argResolvedType = arg?.value?.$resolvedType;
        if (!argResolvedType) {
            accept('error', 'argument type cannot be resolved', { node: arg });
            return false;
        }

        const dstType = param.type.type;
        if (!dstType) {
            accept('error', 'parameter type cannot be resolved', { node: param });
            return false;
        }

        const dstIsArray = param.type.array;
        const dstRef = param.type.reference;

        if (dstType === 'Any' && !dstIsArray) {
            // scalar 'any' can be assigned with anything
            return true;
        }

        if (typeof argResolvedType.decl === 'string') {
            // scalar type
            if (!typeAssignable(dstType, argResolvedType.decl, arg.value) || dstIsArray !== argResolvedType.array) {
                accept('error', `argument is not assignable to parameter`, {
                    node: arg,
                });
                return false;
            }
        } else {
            // enum or model type
            if ((dstRef?.ref !== argResolvedType.decl && dstType !== 'Any') || dstIsArray !== argResolvedType.array) {
                accept('error', `argument is not assignable to parameter`, {
                    node: arg,
                });
                return false;
            }
        }

        return true;
    }

    @func('check')
    private _checkCheck(expr: InvocationExpr, accept: ValidationAcceptor) {
        let valid = true;

        const fieldArg = expr.args[0].value;
        if (!isDataModelFieldReference(fieldArg) || !isDataModel(fieldArg.$resolvedType?.decl)) {
            accept('error', 'argument must be a relation field', { node: expr.args[0] });
            valid = false;
        }

        if (fieldArg.$resolvedType?.array) {
            accept('error', 'argument cannot be an array field', { node: expr.args[0] });
            valid = false;
        }

        const opArg = expr.args[1]?.value;
        if (opArg) {
            const operation = getLiteral<string>(opArg);
            if (!operation || !['read', 'create', 'update', 'delete'].includes(operation)) {
                accept('error', 'argument must be a "read", "create", "update", or "delete"', { node: expr.args[1] });
                valid = false;
            }
        }

        if (!valid) {
            return;
        }

        // check for cyclic relation checking
        const start = fieldArg.$resolvedType?.decl as DataModel;
        const tasks = [expr];
        const seen = new Set<DataModel>();

        while (tasks.length > 0) {
            const currExpr = tasks.pop()!;
            const arg = currExpr.args[0]?.value;

            if (!isDataModel(arg?.$resolvedType?.decl)) {
                continue;
            }

            const currModel = arg.$resolvedType.decl;

            if (seen.has(currModel)) {
                if (currModel === start) {
                    accept('error', 'cyclic dependency detected when following the `check()` call', { node: expr });
                } else {
                    // a cycle is detected but it doesn't start from the invocation expression we're checking,
                    // just break here and the cycle will be reported when we validate the start of it
                }
                break;
            } else {
                seen.add(currModel);
            }

            const policyAttrs = currModel.attributes.filter(
                (attr: DataModelAttribute) => attr.decl.$refText === '@@allow' || attr.decl.$refText === '@@deny'
            );
            for (const attr of policyAttrs) {
                const rule = attr.args[1];
                if (!rule) {
                    continue;
                }
                streamAst(rule).forEach((node) => {
                    if (isCheckInvocation(node)) {
                        tasks.push(node as InvocationExpr);
                    }
                });
            }
        }
    }
}