Skip to content

Instantly share code, notes, and snippets.

@tsandall
Created November 12, 2020 22:40
Show Gist options
  • Save tsandall/f575d4dfbacc175d24d1a1ad5c51fedf to your computer and use it in GitHub Desktop.
Save tsandall/f575d4dfbacc175d24d1a1ad5c51fedf to your computer and use it in GitHub Desktop.
$ go test -bench=. -run=^$ x_bench_test.go
map[object.union:8 sprintf:13]
map[main/main:1 policy/com.styra.kubernetes.validating/main/main:0]
goos: linux
goarch: amd64
BenchmarkEval-20 map[object.union:8 sprintf:13]
map[main/main:1 policy/com.styra.kubernetes.validating/main/main:0]
map[object.union:8 sprintf:13]
map[main/main:1 policy/com.styra.kubernetes.validating/main/main:0]
5 216041540 ns/op
PASS
ok command-line-arguments 2.304s
package main
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io/ioutil"
"reflect"
"strconv"
"testing"
"time"
"github.com/mathetake/gasm/hostfunc"
"github.com/mathetake/gasm/wasm"
"github.com/open-policy-agent/opa/ast"
"github.com/open-policy-agent/opa/metrics"
"github.com/open-policy-agent/opa/topdown"
)
func BenchmarkEval(b *testing.B) {
buf, err := ioutil.ReadFile("policy.wasm")
if err != nil {
panic(err)
}
mod, err := wasm.DecodeModule(bytes.NewBuffer(buf))
if err != nil {
panic(err)
}
builder := hostfunc.NewModuleBuilder()
var vm *wasm.VirtualMachine
builtinFuncs := map[uint64]topdown.BuiltinFunc{}
var builtinIndex map[string]interface{}
builder.MustSetFunction("env", "opa_abort", func(*wasm.VirtualMachine) reflect.Value {
return reflect.ValueOf(func(int32) {
panic("opa_abort")
})
})
builder.MustSetFunction("env", "opa_builtin0", func(*wasm.VirtualMachine) reflect.Value {
return reflect.ValueOf(func(int32, int32) int32 {
panic("opa_builtin0")
})
})
builder.MustSetFunction("env", "opa_builtin1", func(*wasm.VirtualMachine) reflect.Value {
return reflect.ValueOf(func(int32, int32, int32) int32 {
panic("opa_builtin0")
})
})
builder.MustSetFunction("env", "opa_builtin2", func(*wasm.VirtualMachine) reflect.Value {
return reflect.ValueOf(func(id int32, _ int32, a int32, b int32) int32 {
var ax, bx interface{}
opa_json_dump(vm, uint64(a), &ax)
opa_json_dump(vm, uint64(b), &bx)
op1 := ast.MustInterfaceToValue(ax)
op2 := ast.MustInterfaceToValue(bx)
args := []*ast.Term{ast.NewTerm(op1), ast.NewTerm(op2)}
builtinCtx := topdown.BuiltinContext{
Context: context.Background(),
Cancel: nil,
Runtime: nil,
Time: ast.NumberTerm(json.Number(strconv.FormatInt(time.Now().UnixNano(), 10))),
Metrics: metrics.New(),
Location: nil,
Tracers: nil,
QueryID: 0,
ParentID: 0,
}
var result *ast.Term
builtinFuncs[uint64(id)](builtinCtx, args, func(r *ast.Term) error {
result = r
return nil
})
addr := opa_json_parse(vm, []byte(result.String()))
return int32(addr)
})
})
builder.MustSetFunction("env", "opa_builtin3", func(*wasm.VirtualMachine) reflect.Value {
return reflect.ValueOf(func(int32, int32, int32, int32, int32) int32 {
panic("opa_builtin3")
})
})
builder.MustSetFunction("env", "opa_builtin4", func(*wasm.VirtualMachine) reflect.Value {
return reflect.ValueOf(func(int32, int32, int32, int32, int32, int32) int32 {
panic("opa_builtin4")
})
})
extras := builder.Done()
extras["env"].SecExports["memory"] = &wasm.ExportSegment{Name: "memory", Desc: &wasm.ExportDesc{Kind: 0x02, Index: 0}}
extras["env"].IndexSpace.Memory = [][]byte{
make([]byte, 0),
}
mod.SecMemory = []*wasm.LimitsType{&wasm.LimitsType{
Min: 2,
}}
vm, err = wasm.NewVM(mod, extras)
if err != nil {
panic(err)
}
values, _, err := vm.ExecExportedFunction("malloc", 0)
if err != nil {
panic(err)
}
_ = values
values, _, err = vm.ExecExportedFunction("opa_heap_ptr_get")
if err != nil {
panic(err)
}
builtinIndex = builtins(vm)
for name, id := range builtinIndex {
f := topdown.GetBuiltin(name)
if f == nil {
panic(name)
}
n, err := id.(json.Number).Int64()
if err != nil {
panic(err)
}
builtinFuncs[uint64(n)] = f
}
fmt.Println(builtinIndex)
fmt.Println(entrypoints(vm))
b.ResetTimer()
for i := 0; i < b.N; i++ {
result := eval(vm, []byte(`{
"apiVersion": "admission.k8s.io/v1",
"kind": "AdmissionReview",
"request": {
"dryRun": false,
"kind": {
"group": "",
"kind": "Pod",
"version": "v1"
},
"name": "nginx-deny4",
"namespace": "stress-test-2",
"object": {
"apiVersion": "v1",
"kind": "Pod",
"metadata": {
"creationTimestamp": "2020-11-04T21:50:46Z",
"labels": {
"run": "nginx-deny4"
},
"name": "nginx-deny4",
"namespace": "stress-test-2",
"uid": "623ebe52-306a-4a5a-b6b0-799fb657298e"
},
"spec": {
"containers": [
{
"image": "nginx",
"imagePullPolicy": "Always",
"name": "nginx-deny4",
"resources": {},
"terminationMessagePath": "/dev/termination-log",
"terminationMessagePolicy": "File",
"volumeMounts": [
{
"mountPath": "/var/run/secrets/kubernetes.io/serviceaccount",
"name": "default-token-f6r4l",
"readOnly": true
}
]
}
],
"dnsPolicy": "ClusterFirst",
"enableServiceLinks": true,
"priority": 0,
"restartPolicy": "Never",
"schedulerName": "default-scheduler",
"securityContext": {},
"serviceAccount": "default",
"serviceAccountName": "default",
"terminationGracePeriodSeconds": 30,
"tolerations": [
{
"effect": "NoExecute",
"key": "node.kubernetes.io/not-ready",
"operator": "Exists",
"tolerationSeconds": 300
},
{
"effect": "NoExecute",
"key": "node.kubernetes.io/unreachable",
"operator": "Exists",
"tolerationSeconds": 300
}
],
"volumes": [
{
"name": "default-token-f6r4l",
"secret": {
"secretName": "default-token-f6r4l"
}
}
]
},
"status": {
"phase": "Pending",
"qosClass": "BestEffort"
}
},
"oldObject": null,
"operation": "CREATE",
"options": {
"apiVersion": "meta.k8s.io/v1",
"fieldManager": "kubectl-run",
"kind": "CreateOptions"
},
"requestKind": {
"group": "",
"kind": "Pod",
"version": "v1"
},
"requestResource": {
"group": "",
"resource": "pods",
"version": "v1"
},
"resource": {
"group": "",
"resource": "pods",
"version": "v1"
},
"uid": "dca2f108-954e-495b-b519-81cdffb0291a",
"userInfo": {
"groups": [
"system:serviceaccounts",
"system:serviceaccounts:styra-system",
"system:authenticated"
],
"uid": "c218d52b-8a48-11ea-8990-42010aa80067",
"username": "system:serviceaccount:styra-system:default"
}
}
}`))
_ = result
}
// fmt.Println(result)
}
func builtins(vm *wasm.VirtualMachine) map[string]interface{} {
values, _, err := vm.ExecExportedFunction("builtins")
if err != nil {
panic(err)
}
result := map[string]interface{}{}
opa_json_dump(vm, values[0], &result)
return result
}
func entrypoints(vm *wasm.VirtualMachine) map[string]interface{} {
values, _, err := vm.ExecExportedFunction("entrypoints")
if err != nil {
panic(err)
}
result := map[string]interface{}{}
opa_json_dump(vm, values[0], &result)
return result
}
func opa_json_dump(vm *wasm.VirtualMachine, addr uint64, result interface{}) {
raddr := call(vm, "opa_json_dump", addr)
data := vm.Memory[raddr:]
n := bytes.IndexByte(data, 0)
if n < 0 {
n = 0
}
decoder := json.NewDecoder(bytes.NewReader(data[0:n]))
decoder.UseNumber()
if err := decoder.Decode(&result); err != nil {
panic(err)
}
}
func opa_json_parse(vm *wasm.VirtualMachine, a []byte) uint64 {
pos := malloc(vm, uint64(len(a)))
copy(vm.Memory[pos:pos+uint64(len(a))], a)
addr := call(vm, "opa_json_parse", pos, uint64(len(a)))
return addr
}
func malloc(vm *wasm.VirtualMachine, a uint64) uint64 {
return call(vm, "malloc", a)
}
func opa_eval_ctx_new(vm *wasm.VirtualMachine) uint64 {
return call(vm, "opa_eval_ctx_new")
}
func opa_eval_ctx_set_input(vm *wasm.VirtualMachine, a, b uint64) {
callVoid(vm, "opa_eval_ctx_set_input", a, b)
}
func opa_eval_ctx_get_result(vm *wasm.VirtualMachine, a uint64) uint64 {
return call(vm, "opa_eval_ctx_get_result", a)
}
func eval(vm *wasm.VirtualMachine, a []byte) interface{} {
evalCtx := opa_eval_ctx_new(vm)
inputAddr := opa_json_parse(vm, a)
opa_eval_ctx_set_input(vm, evalCtx, inputAddr)
callVoid(vm, "eval", evalCtx)
resultAddr := opa_eval_ctx_get_result(vm, evalCtx)
var result interface{}
opa_json_dump(vm, resultAddr, &result)
return result
}
func call(vm *wasm.VirtualMachine, name string, args ...uint64) uint64 {
values, _, err := vm.ExecExportedFunction(name, args...)
if err != nil {
panic(err)
}
return values[0]
}
func callVoid(vm *wasm.VirtualMachine, name string, args ...uint64) {
values, _, err := vm.ExecExportedFunction(name, args...)
if err != nil {
panic(err)
}
_ = values
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment