Zenn
🔨

OPAの拡張（カスタム関数）

2021/12/22に公開
Golang
Rego
OPA
tech

この記事はOPA/Regoアドベントカレンダーの22日目です。

OPA上で動くRegoは非常に柔軟で、自由な表現ができます。さらにOPA自身もいろいろと拡張ができる仕様となっています。詳しくは公式ドキュメントにて説明がありますが、今回は拡張機能の要の一つであるカスタム関数（Custom Built-in Functions）について紹介したいと思います。

カスタム関数とは

一言でいうと Goで実装した関数をRego内で利用できるようにする 機能です。Regoの仕様を拡張し、ランタイム内で使える関数を増やすという目的で利用します。具体的には 1) 自分の実装に組み込んだランタイムとして使用する、2) カスタム関数を組み込んだ新しい opa バイナリを作成する、という2種類の使い方があります。

OPAは主に以下の2つユースケースを想定してカスタム関数を提供しています。

  • 複雑な処理：暗号系の処理のような複雑な計算を伴うものはGo側で処理するのが良さそうです。例えばJWTの検証は現状すでに組み込み関数として存在しますが、似たような検証のプロセスをRegoでフルスクラッチで記述するのはかなり困難です。そのため、既存のライブラリなどと組み合わせて機能を実装するため、Goによって記述されたロジックを持ち込めるようになっています。
  • I/Oが発生する処理：Regoは「ポリシーの評価」という目的に主眼を置いているためか、原則として[1]外部とのI/O機能は提供されていません。しかしポリシーを評価する際には外部のデータベースを参照するということは発生しうる[2]と考えられ、アクセスのための機能を何らかの形で使いたくなります。また外部のデータベースへのアクセスでは認証認可の処理が必要になるケースも多く、Regoだけでこの機能を実装するのは難しいと考えられます。こちらについても既存のGoのライブラリなどと組み合わせて実装するのが現実的でしょう。

自分の実装に組み込んだランタイムとして使用

まずは自分のGoの実装に組み込む方法について見ていきましょう。公式ドキュメントではHello world的実装が例になっているので、今回は2つの値を比較してそれが異なっていたらどう違っていたかを返り値とする assert という関数を作ってみました。（Regoのテストにあると便利そうです）

package main

import (
	"context"
	"fmt"

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

func main() {
	assertFunc := rego.Function3(
		&rego.Function{
			Name: "assert",
			Decl: types.NewFunction(types.Args(types.S, types.A, types.A), types.S),
		},
		func(_ rego.BuiltinContext, op1, op2, op3 *ast.Term) (*ast.Term, error) {
			if op2.Value.Compare(op3.Value) == 0 {
				return nil, nil
			}

			msg := ast.StringTerm(fmt.Sprintf("Failed '%s': expected %v, but got %v", op1.String(), op2.String(), op3.String()))
			return msg, nil
		},
    )

	r := rego.New(
		rego.Query(`ret := assert("match A is B", "blue", "orange")`),
		assertFunc,
	)

	rs, err := r.Eval(context.Background())
	if err != nil {
		panic(err.Error())
	}

	pp.Println(rs)
}

こちらを実行すると以下のような結果になります。

% go run .
rego.ResultSet{
  rego.Result{
    Expressions: []*rego.ExpressionValue{
      &rego.ExpressionValue{
        Value:    true,
        Text:     "ret := assert(\"match A is B\", \"blue\", \"orange\")",
        Location: &rego.Location{
          Row: 1,
          Col: 1,
        },
      },
    },
    Bindings: rego.Vars{
      "ret": "Failed 'match A is B': expected \"blue\", but got \"orange\"",
    },
  },
}

assert は1番目が「何を比較したか、どうあるべきか」、2番目が「期待される値」、3番目が「実際の値」という、よくあるテスト用の関数をイメージしてもらえればと思います。もし２番目と３話目が違ったらメッセージを返します。例では assert("match A is B", "blue", "orange") という関数が実行され、値が異なったため Failed 'match A is B': expected "blue", but got "orange" という結果が ret に代入されています。

「なんでこんな面倒くさそうなことするの？」と思った方は、ぜひOPAの記述例の回で挙げたTable Driven Testについて参照いただければと思います。

それでは、順番にコードの解説をしていきたいと思います。

	assertFunc := rego.Function3(

rego.Function3 が作成するための関数です。3 というsuffixは3つの引数をとる関数という意味になります[3]。この関数は第1引数が作成するカスタム関数の定義、第2引数がカスタム関数の処理コードになります。

    &rego.Function{
        Name: "assert",
        Decl: types.NewFunction(types.Args(types.S, types.A, types.A), types.S),
    },

まずカスタム関数に関する宣言です。ここではカスタム関数の名前、および引数と返り値の型を定義しています。types.NewFunction が引数と返り値を定義しており、types.Args(types.S, types.A, types.A) で引数が「String」「Any」「Any」であることを示しており、二番目の types.S が「String」の返り値を示しています。

    func(_ rego.BuiltinContext, op1, op2, op3 *ast.Term) (*ast.Term, error) {

これがカスタム関数の実際の処理をするコールバック関数です。今回はI/Oのようなキャンセルやタイムアウトを気にする処理はないため、 context.Context 互換の第一引数は無視しています。第2、3、4引数がカスタム関数に渡された値あるいは変数を表しています。返り値は *ast.Term を nil で返せばカスタム関数が偽だったことになり、評価が失敗します。一方、error を返すと処理が異常終了したことを意味し、評価が中断されます。

    if op2.Value.Compare(op3.Value) == 0 {
        return nil, nil
    }

ここがまずカスタム関数の第2引数と第3引数を比較している部分です。Compare は同値とみなした場合 0 が返されるので、その場合は nil 、つまり関数としての評価は失敗したことを返します。

    msg := ast.StringTerm(fmt.Sprintf("Failed '%s': expected %v, but got %v", op1.String(), op2.String(), op3.String()))
    return msg, nil

最後に msg を文字列型の句として作成し、返り値として渡します。これによってこのカスタム関数の評価が成立し、fail[msg] のような代入をするための変数を作成することができます。

	r := rego.New(
		rego.Query(`ret := assert("match A is B", "blue", "orange")`),
		assertFunc,
	)

最後にRegoのランタイム作成時に、このカスタム関数を埋め込むことで、独自の関数の実行が可能になります。

カスタム関数を組み込んだ新しいバイナリの作成

さて、これで無事にカスタム関数が使えるようになったなったのですが実運用上ではいささか問題があります。関数を埋め込んだコードは特定の目的で動かすには適していますが、既存の opa コマンドからは不正なコードと判定されるようになってしまいます。そのため check が通らなかったり、deps で依存関係の分析ができなかったり、なにより test によってポリシーのテストができません。

そこでこの問題を解決する方法として、 カスタム関数を組み込んだ新しいopaバイナリを作る 手段が提供されています。コードも先程のカスタム関数＋ランタイムと大部分が似通っており、コードを整理すれば独自のランタイムとカスタム関数組み込みバイナリを同じリポジトリで管理するのも容易です。

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"
)

func main() {
	rego.RegisterBuiltin3(
		&rego.Function{
			Name: "assert",
			Decl: types.NewFunction(types.Args(types.S, types.A, types.A), types.S),
		},
		func(_ rego.BuiltinContext, op1, op2, op3 *ast.Term) (*ast.Term, error) {
			if op2.Value.Compare(op3.Value) == 0 {
				return nil, nil
			}

			msg := ast.StringTerm(fmt.Sprintf("Failed '%s': expected %v, but got %v", op1.String(), op2.String(), op3.String()))
			return msg, nil
		},
	)

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

これをビルドすると、

$ go build -o myopa .
$ ./myopa
An open source project to policy-enable your service.

Usage:
  myopa [command]

Available Commands:
  bench       Benchmark a Rego query
  build       Build an OPA bundle
（割愛）

という感じで、実際の opa コマンドと同様の機能＋カスタム関数のバイナリが作成できてしまいます。見ての通り rego.RegisterBuiltin3 がカスタム関数を組み込む役割を果たしており、中身は先程の rego.Function3 と全くおなじになっています。このコードから作成したバイナリを使って先程と同じクエリを発行すると同等の結果を得ることができます。

$ ./myopa eval 'ret := assert("match A is B", "blue", "orange")'
{
  "result": [
    {
      "expressions": [
        {
          "value": true,
          "text": "ret := assert(\"match A is B\", \"blue\", \"orange\")",
          "location": {
            "row": 1,
            "col": 1
          }
        }
      ],
      "bindings": {
        "ret": "Failed '\"match A is B\"': expected \"blue\", but got \"orange\""
      }
    }
  ]
}

まとめ

OPAはとても開かれた形で機能が提供されており、このようなカスタマイズの自由度も高く利用することができます。もちろん実際の運用において自分たちでバイナリを管理するのは一定負荷がかかるため、どのような形態が望ましいかというベストプラクティスはまだこれから醸成されていくと思います。とはいえ主体的に選択できる余地があることで、より自分たちのユースケースにあったOPAの使い方ができるのではないかと考えられます。

脚注

  1. 現状だと唯一シンプルなHTTPリクエスト送信の組み込み関数は用意されています https://www.openpolicyagent.org/docs/latest/policy-reference/#http ↩︎

  2. allow・denyリストなどを扱うために data (basic document) の仕組みはありますが、リアルタイムな更新を考えると外部のデータベース参照も視野に入れて検討するべきと考えられます ↩︎

  3. 引数の固定長は1〜4まであり、それ以外は FunctionDyn を使います ↩︎

Discussion