From 09acbbdc7f6b9284c75342838b9a2b61463c32d8 Mon Sep 17 00:00:00 2001
From: Torin Sandall <torinsandall@gmail.com>
Date: Fri, 5 Oct 2018 19:36:51 -0700
Subject: [PATCH] Add build subcommand and expose compiler

These changes add a build command that writes out a policy.wasm file by
compiling an input query to WASM. These changes also add a simple
compile interface to the rego package.

Signed-off-by: Torin Sandall <torinsandall@gmail.com>
---
 .gitignore   |  1 +
 cmd/build.go | 94 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 rego/rego.go | 46 +++++++++++++++++++++++++
 3 files changed, 141 insertions(+)
 create mode 100644 cmd/build.go

diff --git a/.gitignore b/.gitignore
index b633705895..cff66bae19 100644
--- a/.gitignore
+++ b/.gitignore
@@ -9,6 +9,7 @@ opa_*
 .Dockerfile_*
 _release
 site.tar.gz
+policy.wasm
 
 # runtime artifacts
 policies
diff --git a/cmd/build.go b/cmd/build.go
new file mode 100644
index 0000000000..c3f14b5c18
--- /dev/null
+++ b/cmd/build.go
@@ -0,0 +1,94 @@
+// Copyright 2018 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+package cmd
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"github.com/open-policy-agent/opa/loader"
+	"github.com/open-policy-agent/opa/storage/inmem"
+
+	"github.com/open-policy-agent/opa/rego"
+	"github.com/spf13/cobra"
+)
+
+var buildParams = struct {
+	outputFile string
+	debug      bool
+	dataPaths  repeatedStringFlag
+	ignore     []string
+}{}
+
+var buildCommand = &cobra.Command{
+	Use:   "build <query>",
+	Short: "Compile Rego policy queries",
+	Long: `Compile a Rego policy query into an executable for enforcement.
+
+The 'build' command takes a policy query as input and compiles it into an
+executable that can be loaded into an enforcement point and evaluated with
+input values. By default, the build command produces WebAssembly (WASM)
+executables.`,
+	PreRunE: func(Cmd *cobra.Command, args []string) error {
+		if len(args) == 0 {
+			return fmt.Errorf("specify query argument")
+		}
+		return nil
+	},
+	Run: func(cmd *cobra.Command, args []string) {
+		if err := build(args); err != nil {
+			fmt.Println("error:", err)
+			os.Exit(1)
+		}
+	},
+}
+
+func build(args []string) error {
+
+	ctx := context.Background()
+
+	f := loaderFilter{
+		Ignore: buildParams.ignore,
+	}
+
+	regoArgs := []func(*rego.Rego){
+		rego.Query(args[0]),
+	}
+
+	loaded, err := loader.Filtered(buildParams.dataPaths.v, f.Apply)
+	if err != nil {
+		return err
+	}
+
+	regoArgs = append(regoArgs, rego.Store(inmem.NewFromObject(loaded.Documents)))
+	for _, file := range loaded.Modules {
+		regoArgs = append(regoArgs, rego.Module(file.Name, string(file.Raw)))
+	}
+
+	r := rego.New(regoArgs...)
+	cr, err := r.Compile(ctx)
+	if err != nil {
+		return err
+	}
+
+	out, err := os.Create(buildParams.outputFile)
+	if err != nil {
+		return err
+	}
+
+	defer out.Close()
+
+	_, err = out.Write(cr.Bytes)
+	return err
+}
+
+func init() {
+	buildCommand.Flags().StringVarP(&buildParams.outputFile, "output", "o", "policy.wasm", "set the filename of the compiled policy")
+	buildCommand.Flags().BoolVarP(&buildParams.debug, "debug", "D", false, "enable debug output")
+	buildCommand.Flags().VarP(&buildParams.dataPaths, "data", "d", "set data file(s) or directory path(s)")
+	setIgnore(buildCommand.Flags(), &buildParams.ignore)
+	RootCommand.AddCommand(buildCommand)
+}
diff --git a/rego/rego.go b/rego/rego.go
index 161dbb7232..e7ab9384b8 100644
--- a/rego/rego.go
+++ b/rego/rego.go
@@ -6,10 +6,15 @@
 package rego
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"strings"
 
+	"github.com/open-policy-agent/opa/internal/compiler/wasm"
+	"github.com/open-policy-agent/opa/internal/planner"
+	"github.com/open-policy-agent/opa/internal/wasm/encoding"
+
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/metrics"
 	"github.com/open-policy-agent/opa/storage"
@@ -20,6 +25,12 @@ import (
 
 const defaultPartialNamespace = "partial"
 
+// CompileResult represents sthe result of compiling a Rego query, zero or more
+// Rego modules, and arbitrary contextual data into an executable.
+type CompileResult struct {
+	Bytes []byte `json:"bytes"`
+}
+
 // PartialQueries contains the queries and support modules produced by partial
 // evaluation.
 type PartialQueries struct {
@@ -442,6 +453,41 @@ func (r *Rego) Partial(ctx context.Context) (*PartialQueries, error) {
 	return r.partial(ctx, compiled, txn, partialNamespace)
 }
 
+// Compile returnss a compiled policy query.
+func (r *Rego) Compile(ctx context.Context) (*CompileResult, error) {
+
+	pq, err := r.Partial(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(pq.Support) > 0 {
+		return nil, fmt.Errorf("modules not supported")
+	}
+
+	policy, err := planner.New().WithQueries(pq.Queries).Plan()
+	if err != nil {
+		return nil, err
+	}
+
+	m, err := wasm.New().WithPolicy(policy).Compile()
+	if err != nil {
+		return nil, err
+	}
+
+	var out bytes.Buffer
+
+	if err := encoding.WriteModule(&out, m); err != nil {
+		return nil, err
+	}
+
+	result := &CompileResult{
+		Bytes: out.Bytes(),
+	}
+
+	return result, nil
+}
+
 func (r *Rego) parse() (map[string]*ast.Module, ast.Body, error) {
 
 	r.metrics.Timer(metrics.RegoQueryParse).Start()