config: Allow configuration of capabilities

pkg/config/config.go

@@ -9,12 +9,97 @@ 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"`
+				File    string `yaml:"file"`
+			} `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
+
+	capabilitiesFile := result.Capabilities.From.File
+	capabilitiesEngine := result.Capabilities.From.Engine
+	capabilitiesEngineVersion := result.Capabilities.From.Version
+
+	if capabilitiesFile != "" && capabilitiesEngine != "" {
+		return fmt.Errorf("capabilities from.file and from.engine are mutually exclusive")
+	}
+
+	if capabilitiesEngine != "" && capabilitiesEngineVersion == "" {
+		return fmt.Errorf("please set the version for the engine from which to load capabilities from")
+	}
+
+	if capabilitiesFile != "" {
+		bs, err := os.ReadFile(capabilitiesFile)
+		if err != nil {
+			return fmt.Errorf("failed to load capabilities file: %w", err)
+		}
+
+		err = json.Unmarshal(bs, &config.Capabilities)
+		if err != nil {
+			return fmt.Errorf("failed to unmarshal capabilities file contents: %w", err)
+		}
+	}
+
+	if capabilitiesEngine != "" && 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
+	}
+
+	// by default, use the capabilities from the current OPA
+	if capabilitiesEngine == "" && capabilitiesFile == "" {
+		config.Capabilities = *ast.CapabilitiesForThisVersion()
+	}
+
+	// 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

pkg/config/config_test.go

@@ -7,6 +7,7 @@ import (
 
 	"gopkg.in/yaml.v3"
 
+	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/util/test"
 
 	rio "github.com/styrainc/regal/internal/io"
@@ -123,6 +124,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 +175,120 @@ 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)
+		}
+	}
+}
+
+func TestUnmarshalConfigWithBuiltinsFile(t *testing.T) {
+	t.Parallel()
+
+	bs := []byte(`rules: {}
+capabilities:
+  from:
+    file: "./fixtures/caps.json"
+`)
+
+	var conf Config
+
+	if err := yaml.Unmarshal(bs, &conf); err != nil {
+		t.Fatal(err)
+	}
+
+	if exp, got := 1, len(conf.Capabilities.Builtins); exp != got {
+		t.Errorf("expected %d builtins, got %d", exp, got)
+	}
+
+	expectedBuiltins := []string{"wow"}
+
+	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)
+		}
+	}
+}
+
+func TestUnmarshalConfigDefaultCapabilities(t *testing.T) {
+	t.Parallel()
+
+	bs := []byte(`rules: {}
+`)
+
+	var conf Config
+
+	if err := yaml.Unmarshal(bs, &conf); err != nil {
+		t.Fatal(err)
+	}
+
+	caps := ast.CapabilitiesForThisVersion()
+
+	if exp, got := len(caps.Builtins), len(conf.Capabilities.Builtins); exp != got {
+		t.Errorf("expected %d builtins, got %d", exp, got)
+	}
+
+	// choose the first built-ins to check for to keep the test fast
+	expectedBuiltins := []string{caps.Builtins[0].Name, caps.Builtins[1].Name}
+
+	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)
+		}
+	}
 }

pkg/config/fixtures/caps.json

@@ -0,0 +1,35 @@
+{
+  "builtins": [
+    {
+      "name": "wow",
+      "description": "Increases wow in Rego rule",
+      "categories": [
+        "special"
+      ],
+      "decl": {
+        "args": [
+          {
+            "name": "x",
+            "type": "number"
+          }
+        ],
+        "result": {
+          "description": "the wowness level for input `x`",
+          "name": "y",
+          "type": "number"
+        },
+        "type": "function"
+      }
+    }
+  ],
+  "future_keywords": [
+    "contains",
+    "every",
+    "if",
+    "in"
+  ],
+  "wasm_abi_versions": null,
+  "features": [
+    "rule_head_ref_string_prefixes"
+  ]
+}