通过编写自定义内置函数的方式扩展 OPA/Rego 运行时

前言

本文记录一下如何通过编写自定义内置函数的函数扩展 OPA/Rego 运行时，使得在编写 OPA/Rego 策略语言的时候 可以直接使用自定义的这个内置函数实现更复杂的策略需求。

编写自定义内置函数

我们将要编写的一个自定义内置函数名为 auth.get_user_info ，这个函数的作用是，获取自定义 uid 的用户信息：

user := auth.get_user_info(uid)

并且将使用 github.com/open-policy-agent/opa@v0.35.0 这个 opa 运行时库来编写这个自定义函数， 我们可以使用这个库提供的 rego.RegisterBuiltin1 函数来实现我们的需求：

package main

import (
        "fmt"
        "os"

        "github.com/open-policy-agent/opa/ast"
        "github.com/open-policy-agent/opa/cmd"
        "github.com/open-policy-agent/opa/rego"
        "github.com/open-policy-agent/opa/types"
)

type User struct {
        Name string
        Age int
}

func main() {
        users := map[string]User{
                "uid-1": {
                        Name: "tom",
                        Age:  25,
                },
                "uid-2": {
                        Name: "eric",
                        Age:  30,
                },
        }
        rego.RegisterBuiltin1(
                &rego.Function{
                        Name:    "auth.get_user_info",
                        Decl:    types.NewFunction(types.Args(types.S), types.A),
                        Memoize: true,
                },
                func(bctx rego.BuiltinContext, op1 *ast.Term) (*ast.Term, error) {
                        var uid string
                        if err := ast.As(op1.Value, &uid); err != nil {
                                return nil, err
                        }
                        user, ok := users[uid]
                        if !ok {
                                return nil, fmt.Errorf("user %s is not found", uid)
                        }
                        v, err := ast.InterfaceToValue(user)
                        if err != nil {
                                return nil, err
                        }
                        return ast.NewTerm(v), nil
                },
        )

        if err := cmd.RootCommand.Execute(); err != nil {
                fmt.Println(err)
                os.Exit(1)
        }
}

测试：

$ go build -o custom-built-functions

$ ./custom-built-functions run
OPA 0.35.0 (commit , built at )

Run 'help' to see a list of commands and check for updates.

> uid := "uid-1"
Rule 'uid' defined in package repl. Type 'show' to see rules.
> uid
"uid-1"
> user := auth.get_user_info(uid)
Rule 'user' defined in package repl. Type 'show' to see rules.
> user
{
  "Age": 25,
  "Name": "tom"
}
> user.Age
25
> user.Name
"tom"
>
Do you want to exit ([y]/n)? y

通过上面的测试可以看到，我们已经实现了自定义内置函数的需求，并且测试结果也符合预期结果。

通过上面的示例代码中的 rego.RegisterBuiltin1 这个函数名称可能已经猜到了，如果要定义接受两个参数 的函数的话应该使用 rego.RegisterBuiltin2 函数，同理还有 rego.RegisterBuiltin3 和 reg.RegisterBuiltin4 函数 可供使用。以及还有一个定义不定长参数的 rego.RegisterBuiltinDyn 可以用来满足跟复杂的函数需求。

Comments