Skip to content

Commit

Permalink
feat: support customBuiltins parameter in loadPolicy (#165)
Browse files Browse the repository at this point in the history
This adds support for custom builtins by adding an additional parameter to loadPolicy, somewhat related to the discussion in #23.

One potential point of contention is that this always favors first-party builtins over provided builtins.

Signed-off-by: Adam Berger <[email protected]>
  • Loading branch information
abrgr committed Mar 30, 2022
1 parent ff786cf commit 230febc
Show file tree
Hide file tree
Showing 4 changed files with 271 additions and 8 deletions.
41 changes: 33 additions & 8 deletions src/opa.js
Original file line number Diff line number Diff line change
Expand Up @@ -101,11 +101,18 @@ const builtinFuncs = builtIns;
* @param {WebAssembly.Instance} wasmInstance
* @param {WebAssembly.Memory} memory
* @param {{ [builtinId: number]: string }} builtins
* @param {{ [builtinName: string]: Function }} customBuiltins
* @param {string} builtin_id
*/
function _builtinCall(wasmInstance, memory, builtins, builtinId) {
function _builtinCall(
wasmInstance,
memory,
builtins,
customBuiltins,
builtinId,
) {
const builtInName = builtins[builtinId];
const impl = builtinFuncs[builtInName];
const impl = builtinFuncs[builtInName] || customBuiltins[builtInName];

if (impl === undefined) {
throw {
Expand All @@ -119,7 +126,7 @@ function _builtinCall(wasmInstance, memory, builtins, builtinId) {
const argArray = Array.prototype.slice.apply(arguments);
const args = [];

for (let i = 4; i < argArray.length; i++) {
for (let i = 5; i < argArray.length; i++) {
const jsArg = _dumpJSON(wasmInstance, memory, argArray[i]);
args.push(jsArg);
}
Expand All @@ -131,15 +138,18 @@ function _builtinCall(wasmInstance, memory, builtins, builtinId) {

/**
* _loadPolicy can take in either an ArrayBuffer or WebAssembly.Module
* as its first argument, and a WebAssembly.Memory for the second parameter.
* as its first argument, a WebAssembly.Memory for the second parameter,
* and an object mapping string names to additional builtin functions for
* the third parameter.
* It will return a Promise, depending on the input type the promise
* resolves to both a compiled WebAssembly.Module and its first WebAssembly.Instance
* or to the WebAssemblyInstance.
* @param {BufferSource | WebAssembly.Module} policyWasm
* @param {WebAssembly.Memory} memory
* @param {{ [builtinName: string]: Function }} customBuiltins
* @returns {Promise<{ policy: WebAssembly.WebAssemblyInstantiatedSource | WebAssembly.Instance, minorVersion: number }>}
*/
async function _loadPolicy(policyWasm, memory) {
async function _loadPolicy(policyWasm, memory, customBuiltins) {
const addr2string = stringDecoder(memory);

const env = {};
Expand All @@ -154,13 +164,20 @@ async function _loadPolicy(policyWasm, memory) {
console.log(addr2string(addr));
},
opa_builtin0: function (builtinId, _ctx) {
return _builtinCall(env.instance, memory, env.builtins, builtinId);
return _builtinCall(
env.instance,
memory,
env.builtins,
customBuiltins,
builtinId,
);
},
opa_builtin1: function (builtinId, _ctx, arg1) {
return _builtinCall(
env.instance,
memory,
env.builtins,
customBuiltins,
builtinId,
arg1,
);
Expand All @@ -170,6 +187,7 @@ async function _loadPolicy(policyWasm, memory) {
env.instance,
memory,
env.builtins,
customBuiltins,
builtinId,
arg1,
arg2,
Expand All @@ -180,6 +198,7 @@ async function _loadPolicy(policyWasm, memory) {
env.instance,
memory,
env.builtins,
customBuiltins,
builtinId,
arg1,
arg2,
Expand All @@ -191,6 +210,7 @@ async function _loadPolicy(policyWasm, memory) {
env.instance,
memory,
env.builtins,
customBuiltins,
builtinId,
arg1,
arg2,
Expand Down Expand Up @@ -399,16 +419,21 @@ module.exports = {
* Defaults to 5 pages (320KB).
* @param {BufferSource | WebAssembly.Module} regoWasm
* @param {number | WebAssembly.MemoryDescriptor} memoryDescriptor For backwards-compatibility, a 'number' argument is taken to be the initial memory size.
* @param {{ [builtinName: string]: Function }} customBuiltins A map from string names to builtin functions
*/
async loadPolicy(regoWasm, memoryDescriptor = {}) {
async loadPolicy(regoWasm, memoryDescriptor = {}, customBuiltins = {}) {
// back-compat, second arg used to be a number: 'memorySize', with default of 5
if (typeof memoryDescriptor === "number") {
memoryDescriptor = { initial: memoryDescriptor };
}
memoryDescriptor.initial = memoryDescriptor.initial || 5;

const memory = new WebAssembly.Memory(memoryDescriptor);
const { policy, minorVersion } = await _loadPolicy(regoWasm, memory);
const { policy, minorVersion } = await _loadPolicy(
regoWasm,
memory,
customBuiltins,
);
return new LoadedPolicy(policy, memory, minorVersion);
},
};
98 changes: 98 additions & 0 deletions test/fixtures/custom-builtins/capabilities.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,98 @@
{
"builtins": [
{
"name": "custom.zeroArgBuiltin",
"decl": {
"type": "function",
"args": [],
"result": {
"type": "string"
}
}
},
{
"name": "custom.oneArgBuiltin",
"decl": {
"type": "function",
"args": [
{ "type": "string" }
],
"result": {
"type": "string"
}
}
},
{
"name": "custom.twoArgBuiltin",
"decl": {
"type": "function",
"args": [
{ "type": "string" },
{ "type": "string" }
],
"result": {
"type": "string"
}
}
},
{
"name": "custom.threeArgBuiltin",
"decl": {
"type": "function",
"args": [
{ "type": "string" },
{ "type": "string" },
{ "type": "string" }
],
"result": {
"type": "string"
}
}
},
{
"name": "custom.fourArgBuiltin",
"decl": {
"type": "function",
"args": [
{ "type": "string" },
{ "type": "string" },
{ "type": "string" },
{ "type": "string" }
],
"result": {
"type": "string"
}
}
},
{
"name": "json.is_valid",
"decl": {
"type": "function",
"args": [
{ "type": "string" }
],
"result": {
"type": "boolean"
}
}
},
{
"name": "eq",
"decl": {
"args": [
{
"type": "any"
},
{
"type": "any"
}
],
"result": {
"type": "boolean"
},
"type": "function"
},
"infix": "="
}
]
}
25 changes: 25 additions & 0 deletions test/fixtures/custom-builtins/custom-builtins-policy.rego
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
package custom_builtins

zero_arg = x {
x = custom.zeroArgBuiltin()
}

one_arg = x {
x = custom.oneArgBuiltin(input.args[0])
}

two_arg = x {
x = custom.twoArgBuiltin(input.args[0], input.args[1])
}

three_arg = x {
x = custom.threeArgBuiltin(input.args[0], input.args[1], input.args[2])
}

four_arg = x {
x = custom.fourArgBuiltin(input.args[0], input.args[1], input.args[2], input.args[3])
}

valid_json {
json.is_valid("{}")
}
115 changes: 115 additions & 0 deletions test/opa-custom-builtins.test.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,115 @@
const { readFileSync } = require("fs");
const { execFileSync } = require("child_process");
const { loadPolicy } = require("../src/opa.js");

describe("custom builtins", () => {
const fixturesFolder = "test/fixtures/custom-builtins";

let policy;

beforeAll(async () => {
const bundlePath = `${fixturesFolder}/bundle.tar.gz`;

execFileSync("opa", [
"build",
fixturesFolder,
"-o",
bundlePath,
"-t",
"wasm",
"--capabilities",
`${fixturesFolder}/capabilities.json`,
"-e",
"custom_builtins/zero_arg",
"-e",
"custom_builtins/one_arg",
"-e",
"custom_builtins/two_arg",
"-e",
"custom_builtins/three_arg",
"-e",
"custom_builtins/four_arg",
"-e",
"custom_builtins/valid_json",
]);

execFileSync("tar", [
"-xzf",
bundlePath,
"-C",
`${fixturesFolder}/`,
"/policy.wasm",
]);

const policyWasm = readFileSync(`${fixturesFolder}/policy.wasm`);
const opts = { initial: 5, maximum: 10 };
policy = await loadPolicy(policyWasm, opts, {
"custom.zeroArgBuiltin": () => `hello`,
"custom.oneArgBuiltin": (arg0) => `hello ${arg0}`,
"custom.twoArgBuiltin": (arg0, arg1) => `hello ${arg0}, ${arg1}`,
"custom.threeArgBuiltin": (arg0, arg1, arg2) => (
`hello ${arg0}, ${arg1}, ${arg2}`
),
"custom.fourArgBuiltin": (arg0, arg1, arg2, arg3) => (
`hello ${arg0}, ${arg1}, ${arg2}, ${arg3}`
),
"json.is_valid": () => {
throw new Error("should never happen");
},
});
});

it("should call a custom zero-arg builtin", () => {
const result = policy.evaluate({}, "custom_builtins/zero_arg");
expect(result.length).not.toBe(0);
expect(result[0]).toMatchObject({ result: "hello" });
});

it("should call a custom one-arg builtin", () => {
const result = policy.evaluate(
{ args: ["arg0"] },
"custom_builtins/one_arg",
);
expect(result.length).not.toBe(0);
expect(result[0]).toMatchObject({ result: "hello arg0" });
});

it("should call a custom two-arg builtin", () => {
const result = policy.evaluate(
{ args: ["arg0", "arg1"] },
"custom_builtins/two_arg",
);
expect(result.length).not.toBe(0);
expect(result[0]).toMatchObject({
result: "hello arg0, arg1",
});
});

it("should call a custom three-arg builtin", () => {
const result = policy.evaluate(
{ args: ["arg0", "arg1", "arg2"] },
"custom_builtins/three_arg",
);
expect(result.length).not.toBe(0);
expect(result[0]).toMatchObject({
result: "hello arg0, arg1, arg2",
});
});

it("should call a custom four-arg builtin", () => {
const result = policy.evaluate(
{ args: ["arg0", "arg1", "arg2", "arg3"] },
"custom_builtins/four_arg",
);
expect(result.length).not.toBe(0);
expect(result[0]).toMatchObject({
result: "hello arg0, arg1, arg2, arg3",
});
});

it("should call a provided builtin over a custom builtin", () => {
const result = policy.evaluate({}, "custom_builtins/valid_json");
expect(result.length).not.toBe(0);
expect(result[0]).toMatchObject({ result: true });
});
});

0 comments on commit 230febc

Please sign in to comment.