Add support for embedding custom Rego rules (library and CLI) (#355)

* Add support for embedding custom Rego rules in Poutine library

This enhancement allows library consumers (like pkg-supply and spicy-poutine) to embed their own custom Rego rules directly into their binaries alongside Poutine's built-in rules, creating fully self-contained deployments without filesystem dependencies.

Changes:
- Add NewOpaWithEmbeddedRules() constructor that accepts embed.FS containing custom rules
- Add AddEmbeddedRules() method for adding rules to existing Opa instances
- Modify Compile() to load custom embedded rules alongside built-in rules
- Custom rules respect skip and allowed filters like filesystem-based rules
- Fully backward compatible with existing NewOpa() usage

Usage example:
  //go:embed rules/*.rego
  var CustomRules embed.FS

  opa, err := poutineOpa.NewOpaWithEmbeddedRules(ctx, config, CustomRules, "rules")

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* Fix wrapcheck lint issues in custom embedded rules loading

Wrap errors from embed.FS.ReadFile() and fs.WalkDir() with context
to satisfy wrapcheck linter for new code.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* Add CLI support for custom embedded rules and remove AddEmbeddedRules

- Remove AddEmbeddedRules() method (no clear use case)
- Add CustomEmbeddedRules and CustomEmbeddedRulesRoot exported variables to cmd package
- Update newOpa() and newOpaWithConfig() to use NewOpaWithEmbeddedRules when set
- CLI extensions can now set poutineCmd.CustomEmbeddedRules before Execute()

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* Remove customRoot parameter - always use "." as root

The customRoot parameter was unnecessary implementation detail.
Custom embedded rules are now always loaded from "." root.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* Simplify custom embedded rules prefix to "custom/"

Remove unnecessary index from prefix - just use "custom/" like
"poutine/opa/" for built-in and "include/" for filesystem rules.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

---------

Co-authored-by: Claude <[email protected]>

Author: fproulx-boostsecurity

cmd/root.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"context"
+	"embed"
 	"fmt"
 	"os"
 	"os/signal"
@@ -36,6 +37,7 @@ var (
 	Date    string
 )
 var Token string
+var CustomEmbeddedRules *embed.FS
 var cfgFile string
 var config *models.Config = models.DefaultConfig()
 var skipRules []string
@@ -220,7 +222,16 @@ func newOpa(ctx context.Context) (*opa.Opa, error) {
 	if len(allowedRules) > 0 {
 		config.AllowedRules = allowedRules
 	}
-	opaClient, err := opa.NewOpa(ctx, config)
+
+	var opaClient *opa.Opa
+	var err error
+
+	if CustomEmbeddedRules != nil {
+		opaClient, err = opa.NewOpaWithEmbeddedRules(ctx, config, *CustomEmbeddedRules)
+	} else {
+		opaClient, err = opa.NewOpa(ctx, config)
+	}
+
 	if err != nil {
 		log.Error().Err(err).Msg("Failed to create OPA client")
 		return nil, err
@@ -231,7 +242,15 @@ func newOpa(ctx context.Context) (*opa.Opa, error) {
 
 // newOpaWithConfig creates an OPA client with request-scoped configuration
 func newOpaWithConfig(ctx context.Context, cfg *models.Config) (*opa.Opa, error) {
-	opaClient, err := opa.NewOpa(ctx, cfg)
+	var opaClient *opa.Opa
+	var err error
+
+	if CustomEmbeddedRules != nil {
+		opaClient, err = opa.NewOpaWithEmbeddedRules(ctx, cfg, *CustomEmbeddedRules)
+	} else {
+		opaClient, err = opa.NewOpa(ctx, cfg)
+	}
+
 	if err != nil {
 		log.Error().Err(err).Msg("Failed to create OPA client")
 		return nil, fmt.Errorf("failed to create OPA client: %w", err)

opa/opa.go

@@ -27,10 +27,15 @@ var regoFs embed.FS
 //go:embed capabilities.json
 var capabilitiesJson []byte
 
+type embeddedSource struct {
+	fs embed.FS
+}
+
 type Opa struct {
-	Compiler  *ast.Compiler
-	Store     storage.Store
-	LoadPaths []string
+	Compiler            *ast.Compiler
+	Store               storage.Store
+	LoadPaths           []string
+	customEmbeddedRules []embeddedSource
 }
 
 func NewOpa(ctx context.Context, config *models.Config) (*Opa, error) {
@@ -61,6 +66,45 @@ func NewOpa(ctx context.Context, config *models.Config) (*Opa, error) {
 	return newOpa, nil
 }
 
+// NewOpaWithEmbeddedRules creates a new Opa instance with custom embedded Rego rules.
+// This allows library consumers to embed their own Rego rules directly in their binaries
+// alongside Poutine's built-in rules, creating fully self-contained deployments.
+//
+// Example usage:
+//
+//	//go:embed rules
+//	var CustomRules embed.FS
+//
+//	opa, err := poutineOpa.NewOpaWithEmbeddedRules(ctx, config, CustomRules)
+func NewOpaWithEmbeddedRules(ctx context.Context, config *models.Config, customFS embed.FS) (*Opa, error) {
+	registerBuiltinFunctions()
+
+	newOpa := &Opa{
+		Store: inmem.NewFromObject(map[string]interface {
+		}{
+			"config": models.DefaultConfig(),
+		}),
+		customEmbeddedRules: []embeddedSource{{fs: customFS}},
+	}
+
+	if err := newOpa.WithConfig(ctx, config); err != nil {
+		return nil, fmt.Errorf("failed to set opa with config: %w", err)
+	}
+
+	subset := []string{}
+	for _, skip := range config.Skip {
+		if skip.HasOnlyRule() {
+			subset = append(subset, skip.Rule...)
+		}
+	}
+
+	if err := newOpa.Compile(ctx, subset, config.AllowedRules); err != nil {
+		return nil, fmt.Errorf("failed to initialize opa compiler: %w", err)
+	}
+
+	return newOpa, nil
+}
+
 func (o *Opa) Print(ctx print.Context, s string) error {
 	log.Debug().Ctx(ctx.Context).Str("location", ctx.Location.String()).Msg(s)
 	return nil
@@ -103,6 +147,8 @@ func skipRule(path string, skip []string, allowed []string) bool {
 
 func (o *Opa) Compile(ctx context.Context, skip []string, allowed []string) error {
 	modules := make(map[string]string)
+
+	// Load Poutine's built-in rules
 	err := fs.WalkDir(regoFs, "rego", func(path string, d fs.DirEntry, err error) error {
 		if d.IsDir() {
 			return err
@@ -124,6 +170,35 @@ func (o *Opa) Compile(ctx context.Context, skip []string, allowed []string) erro
 		return err
 	}
 
+	// Load custom embedded rules
+	for _, source := range o.customEmbeddedRules {
+		err := fs.WalkDir(source.fs, ".", func(path string, d fs.DirEntry, err error) error {
+			if err != nil {
+				return err
+			}
+
+			if d.IsDir() {
+				return nil
+			}
+
+			if skipRule(path, skip, allowed) {
+				return nil
+			}
+
+			content, err := source.fs.ReadFile(path)
+			if err != nil {
+				return fmt.Errorf("failed to read custom embedded rule %s: %w", path, err)
+			}
+
+			modules["custom/"+path] = string(content)
+			return nil
+		})
+		if err != nil {
+			return fmt.Errorf("failed to load custom embedded rules: %w", err)
+		}
+	}
+
+	// Load rules from filesystem paths
 	result, err := loader.NewFileLoader().
 		WithProcessAnnotation(true).
 		WithRegoVersion(ast.RegoV0CompatV1).

opa/opa_test.go

@@ -2,6 +2,7 @@ package opa
 
 import (
 	"context"
+	"embed"
 	"github.com/boostsecurityio/poutine/models"
 	"github.com/boostsecurityio/poutine/results"
 	"github.com/open-policy-agent/opa/v1/ast"
@@ -11,6 +12,9 @@ import (
 	"testing"
 )
 
+//go:embed testdata/embedded
+var testEmbeddedRules embed.FS
+
 func noOpaErrors(t *testing.T, err error) {
 	if err == nil {
 		return
@@ -250,3 +254,73 @@ func TestWithRulesConfig(t *testing.T) {
 	assert.Equal(t, []interface{}{}, rule.Config["allowed_runners"].Default)
 	assert.Equal(t, []interface{}{"self-hosted"}, rule.Config["allowed_runners"].Value)
 }
+
+func TestNewOpaWithEmbeddedRules(t *testing.T) {
+	ctx := context.TODO()
+
+	// Test NewOpaWithEmbeddedRules constructor
+	opa, err := NewOpaWithEmbeddedRules(ctx, &models.Config{
+		Include: []models.ConfigInclude{},
+	}, testEmbeddedRules)
+	noOpaErrors(t, err)
+	assert.NotNil(t, opa)
+
+	// Verify that the custom rule was loaded and can be evaluated
+	var customRule map[string]interface{}
+	err = opa.Eval(ctx, "data.custom.rule", nil, &customRule)
+	noOpaErrors(t, err)
+	assert.Equal(t, "Custom Test Rule", customRule["title"])
+	assert.Equal(t, "warning", customRule["level"])
+
+	// Test that the custom rule logic works
+	var results []map[string]interface{}
+	input := map[string]interface{}{
+		"test_value": "test data",
+	}
+	err = opa.Eval(ctx, "data.custom.results", input, &results)
+	noOpaErrors(t, err)
+	assert.Len(t, results, 1)
+	assert.Equal(t, "Custom rule executed successfully", results[0]["message"])
+	assert.Equal(t, "test data", results[0]["details"])
+
+	// Verify built-in Poutine rules are still loaded
+	var builtinRule interface{}
+	err = opa.Eval(ctx, "data.rules.pr_runs_on_self_hosted.rule", nil, &builtinRule)
+	noOpaErrors(t, err)
+	assert.NotNil(t, builtinRule)
+}
+
+func TestEmbeddedRulesWithSkipAndAllowed(t *testing.T) {
+	ctx := context.TODO()
+
+	// Test that skip rules work with embedded custom rules
+	opa, err := NewOpaWithEmbeddedRules(ctx, &models.Config{
+		Include: []models.ConfigInclude{},
+	}, testEmbeddedRules)
+	noOpaErrors(t, err)
+
+	// Verify both rules are loaded initially
+	var customRule map[string]interface{}
+	err = opa.Eval(ctx, "data.custom.rule", nil, &customRule)
+	noOpaErrors(t, err)
+	assert.Equal(t, "Custom Test Rule", customRule["title"])
+
+	var skippableRule map[string]interface{}
+	err = opa.Eval(ctx, "data.custom.rules.skippable_rule", nil, &skippableRule)
+	noOpaErrors(t, err)
+	assert.NotNil(t, skippableRule)
+	assert.Equal(t, "Skippable Test Rule", skippableRule["title"])
+
+	// Now recompile with skip rule
+	err = opa.Compile(ctx, []string{"skippable_rule"}, []string{})
+	noOpaErrors(t, err)
+
+	// The non-skipped rule should still be available
+	err = opa.Eval(ctx, "data.custom.rule", nil, &customRule)
+	noOpaErrors(t, err)
+	assert.Equal(t, "Custom Test Rule", customRule["title"])
+
+	// The skipped rule should not be available
+	err = opa.Eval(ctx, "data.custom.rules.skippable_rule", nil, &skippableRule)
+	assert.Error(t, err)
+}

opa/testdata/embedded/custom_rule.rego

@@ -0,0 +1,20 @@
+package custom
+
+# METADATA
+# title: Custom Test Rule
+# description: A custom rule for testing embedded rules functionality
+# custom:
+#   level: warning
+
+rule := {
+	"title": "Custom Test Rule",
+	"description": "This is a custom embedded rule for testing",
+	"level": "warning",
+}
+
+results contains {
+	"message": "Custom rule executed successfully",
+	"details": input.test_value,
+} if {
+	input.test_value != ""
+}

opa/testdata/embedded/rules/skippable_rule.rego

@@ -0,0 +1,13 @@
+package custom.rules
+
+# METADATA
+# title: Skippable Test Rule
+# description: A rule for testing skip functionality
+# custom:
+#   level: warning
+
+skippable_rule := {
+	"title": "Skippable Test Rule",
+	"description": "This rule should be skippable",
+	"level": "warning",
+}