From 089d95ae6e2cbbdd75b50bc00320455c9ee9adc7 Mon Sep 17 00:00:00 2001
From: Charlie Egan <charlie@styra.com>
Date: Thu, 26 Oct 2023 14:38:37 +0100
Subject: [PATCH] config: Allow configuration of capabilities

This is a first pass at making it possible to configure
capabilities in regal.

I haven't yet updated the regoArgs to use this from the config since
I wanted to get some feedback on the impl first.

Signed-off-by: Charlie Egan <charlie@styra.com>
---
 pkg/config/config.go      | 59 +++++++++++++++++++++++++++++++++++++--
 pkg/config/config_test.go | 56 +++++++++++++++++++++++++++++++++++++
 2 files changed, 113 insertions(+), 2 deletions(-)

diff --git a/pkg/config/config.go b/pkg/config/config.go
index 41c6aa363..874df4ab9 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -9,12 +9,67 @@ import (
 
 	"gopkg.in/yaml.v3"
 
+	"github.com/open-policy-agent/opa/ast"
+
 	rio "github.com/styrainc/regal/internal/io"
 )
 
+const capabilitiesEngineOPA = "opa"
+
 type Config struct {
-	Rules  map[string]Category `json:"rules"            yaml:"rules"`
-	Ignore Ignore              `json:"ignore,omitempty" yaml:"ignore,omitempty"`
+	Rules        map[string]Category `json:"rules"            yaml:"rules"`
+	Ignore       Ignore              `json:"ignore,omitempty" yaml:"ignore,omitempty"`
+	Capabilities ast.Capabilities    `json:"capabilities,omitempty"     yaml:"capabilities,omitempty"`
+}
+
+func (config *Config) UnmarshalYAML(value *yaml.Node) error {
+	var result struct {
+		Rules        map[string]Category `yaml:"rules"`
+		Ignore       Ignore              `yaml:"ignore"`
+		Capabilities struct {
+			From struct {
+				Engine  string `yaml:"engine"`
+				Version string `yaml:"version"`
+			} `yaml:"from"`
+			Plus struct {
+				Builtins []*ast.Builtin `yaml:"builtins"`
+			} `yaml:"plus"`
+			Minus struct {
+				Builtins []struct {
+					Name string `yaml:"name"`
+				} `yaml:"builtins"`
+			} `yaml:"minus"`
+		} `yaml:"capabilities"`
+	}
+
+	if err := value.Decode(&result); err != nil {
+		return fmt.Errorf("unmarshalling config failed %w", err)
+	}
+
+	config.Rules = result.Rules
+	config.Ignore = result.Ignore
+
+	if result.Capabilities.From.Engine == capabilitiesEngineOPA {
+		capabilities, err := ast.LoadCapabilitiesVersion(result.Capabilities.From.Version)
+		if err != nil {
+			return fmt.Errorf("loading capabilities failed: %w", err)
+		}
+
+		config.Capabilities = *capabilities
+	}
+
+	// remove any builtins referenced in the minus config
+	for i, builtin := range config.Capabilities.Builtins {
+		for _, minusBuiltin := range result.Capabilities.Minus.Builtins {
+			if minusBuiltin.Name == builtin.Name {
+				config.Capabilities.Builtins = append(config.Capabilities.Builtins[:i], config.Capabilities.Builtins[i+1:]...)
+			}
+		}
+	}
+
+	config.Capabilities.Builtins = append(config.Capabilities.Builtins, result.Capabilities.Plus.Builtins...)
+
+	return nil
 }
 
 type Category map[string]Rule
diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go
index fadf6153a..9ea89d86a 100644
--- a/pkg/config/config_test.go
+++ b/pkg/config/config_test.go
@@ -123,6 +123,22 @@ func TestUnmarshalConfig(t *testing.T) {
         files:
           - foo.rego
       level: error
+capabilities:
+  from:
+    engine: opa
+    version: v0.45.0
+  plus:
+    builtins:
+      - name: ldap.query
+        type: function
+        decl:
+          args:
+            - type: string
+        result:
+          type: object
+  minus:
+    builtins:
+      - name: http.send
 `)
 
 	var conf Config
@@ -158,4 +174,44 @@ func TestUnmarshalConfig(t *testing.T) {
 	if conf.Rules["testing"]["foo"].Extra["level"] != nil {
 		t.Errorf("expected extra attribute 'level' to be removed")
 	}
+
+	if exp, got := 183, len(conf.Capabilities.Builtins); exp != got {
+		t.Errorf("expected %d builtins, got %d", exp, got)
+	}
+
+	expectedBuiltins := []string{"regex.match", "ldap.query"}
+
+	for _, expectedBuiltin := range expectedBuiltins {
+		expectedBuiltinFound := false
+
+		for _, bi := range conf.Capabilities.Builtins {
+			if bi.Name == expectedBuiltin {
+				expectedBuiltinFound = true
+
+				break
+			}
+		}
+
+		if !expectedBuiltinFound {
+			t.Errorf("expected builtin %s to be found", expectedBuiltin)
+		}
+	}
+
+	unexpectedBuiltins := []string{"http.send"}
+
+	for _, unexpectedBuiltin := range unexpectedBuiltins {
+		unexpectedBuiltinFound := false
+
+		for _, bi := range conf.Capabilities.Builtins {
+			if bi.Name == unexpectedBuiltin {
+				unexpectedBuiltinFound = true
+
+				break
+			}
+		}
+
+		if unexpectedBuiltinFound {
+			t.Errorf("expected builtin %s to be removed", unexpectedBuiltin)
+		}
+	}
 }