cli: factor out genpolicy

cli/cmd/common.go

@@ -30,12 +30,6 @@ const (
 )
 
 var (
-	//go:embed assets/genpolicy
-	genpolicyBin []byte
-	//go:embed assets/genpolicy-settings.json
-	defaultGenpolicySettings []byte
-	//go:embed assets/genpolicy-rules.rego
-	defaultRules []byte
 	// ReleaseImageReplacements contains the image replacements used by contrast.
 	//go:embed assets/image-replacements.txt
 	ReleaseImageReplacements []byte

cli/cmd/generate.go

@@ -4,25 +4,20 @@
 package cmd
 
 import (
-	"bufio"
 	"bytes"
 	"context"
 	"crypto/sha256"
-	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io"
 	"io/fs"
 	"log/slog"
 	"os"
-	"os/exec"
 	"path/filepath"
-	"regexp"
 	"slices"
 	"strings"
 
-	"github.com/edgelesssys/contrast/internal/embedbin"
+	"github.com/edgelesssys/contrast/cli/genpolicy"
 	"github.com/edgelesssys/contrast/internal/kuberesource"
 	"github.com/edgelesssys/contrast/internal/manifest"
 	"github.com/edgelesssys/contrast/internal/platforms"
@@ -131,7 +126,7 @@ func runGenerate(cmd *cobra.Command, args []string) error {
 	}
 	fmt.Fprintln(cmd.OutOrStdout(), "✔️ Patched targets")
 
-	if err := generatePolicies(cmd.Context(), flags.policyPath, flags.settingsPath, flags.genpolicyCachePath, paths, log); err != nil {
+	if err := generatePolicies(cmd.Context(), flags, paths, log); err != nil {
 		return fmt.Errorf("generate policies: %w", err)
 	}
 	fmt.Fprintln(cmd.OutOrStdout(), "✔️ Generated workload policy annotations")
@@ -246,36 +241,30 @@ func filterNonCoCoRuntime(runtimeClassNamePrefix string, paths []string, logger
 	return filtered
 }
 
-func generatePolicies(ctx context.Context, regoRulesPath, policySettingsPath, genpolicyCachePath string, yamlPaths []string, logger *slog.Logger) error {
-	if err := createFileWithDefault(policySettingsPath, 0o644, func() ([]byte, error) { return defaultGenpolicySettings, nil }); err != nil {
+func generatePolicies(ctx context.Context, flags *generateFlags, yamlPaths []string, logger *slog.Logger) error {
+	cfg := genpolicy.NewConfig(flags.referenceValuesPlatform)
+	if err := createFileWithDefault(flags.settingsPath, 0o644, func() ([]byte, error) { return cfg.Settings, nil }); err != nil {
 		return fmt.Errorf("creating default policy file: %w", err)
 	}
-	if err := createFileWithDefault(regoRulesPath, 0o644, func() ([]byte, error) { return defaultRules, nil }); err != nil {
+	if err := createFileWithDefault(flags.policyPath, 0o644, func() ([]byte, error) { return cfg.Rules, nil }); err != nil {
 		return fmt.Errorf("creating default policy.rego file: %w", err)
 	}
-	binaryInstallDir, err := installDir()
-	if err != nil {
-		return fmt.Errorf("get install dir: %w", err)
-	}
-	genpolicyInstall, err := embedbin.New().Install(binaryInstallDir, genpolicyBin)
+
+	runner, err := genpolicy.New(flags.policyPath, flags.settingsPath, flags.genpolicyCachePath)
 	if err != nil {
-		return fmt.Errorf("install genpolicy: %w", err)
+		return fmt.Errorf("preparing genpolicy: %w", err)
 	}
+
 	defer func() {
-		if err := genpolicyInstall.Uninstall(); err != nil {
-			logger.Warn("uninstall genpolicy tool", "err", err)
+		if err := runner.Teardown(); err != nil {
+			logger.Warn("Cleanup failed", "err", err)
 		}
 	}()
+
 	for _, yamlPath := range yamlPaths {
-		policyHash, err := generatePolicyForFile(ctx, genpolicyInstall.Path(), regoRulesPath, policySettingsPath, yamlPath, genpolicyCachePath, logger)
-		if err != nil {
-			return fmt.Errorf("generate policy for %s: %w", yamlPath, err)
-		}
-		if policyHash == [32]byte{} {
-			continue
+		if err := runner.Run(ctx, yamlPath, logger); err != nil {
+			return fmt.Errorf("failed to generate policy for %s: %w", yamlPath, err)
 		}
-
-		logger.Info("Calculated policy hash", "hash", hex.EncodeToString(policyHash[:]), "path", yamlPath)
 	}
 	return nil
 }
@@ -419,95 +408,6 @@ func addSeedshareOwnerKeyToManifest(manifst *manifest.Manifest, keyPath string)
 	return nil
 }
 
-type logTranslator struct {
-	r         *io.PipeReader
-	w         *io.PipeWriter
-	logger    *slog.Logger
-	stopDoneC chan struct{}
-}
-
-func newLogTranslator(logger *slog.Logger) logTranslator {
-	r, w := io.Pipe()
-	l := logTranslator{
-		r:         r,
-		w:         w,
-		logger:    logger,
-		stopDoneC: make(chan struct{}),
-	}
-	l.startTranslate()
-	return l
-}
-
-func (l logTranslator) Write(p []byte) (n int, err error) {
-	return l.w.Write(p)
-}
-
-var genpolicyLogPrefixReg = regexp.MustCompile(`^\[[^\]\s]+\s+(\w+)\s+([^\]\s]+)\] (.*)`)
-
-func (l logTranslator) startTranslate() {
-	go func() {
-		defer close(l.stopDoneC)
-		scanner := bufio.NewScanner(l.r)
-		for scanner.Scan() {
-			line := scanner.Text()
-			match := genpolicyLogPrefixReg.FindStringSubmatch(line)
-			if len(match) != 4 {
-				// genpolicy prints some warnings without the logger
-				l.logger.Warn(line)
-			} else {
-				switch match[1] {
-				case "ERROR":
-					l.logger.Error(match[3], "position", match[2])
-				case "WARN":
-					l.logger.Warn(match[3], "position", match[2])
-				case "INFO": // prints quite a lot, only show on debug
-					l.logger.Debug(match[3], "position", match[2])
-				}
-			}
-		}
-	}()
-}
-
-func (l logTranslator) stop() {
-	l.w.Close()
-	<-l.stopDoneC
-}
-
-func generatePolicyForFile(ctx context.Context, genpolicyPath, regoPath, policyPath, yamlPath, genpolicyCachePath string, logger *slog.Logger) ([32]byte, error) {
-	args := []string{
-		"--raw-out",
-		fmt.Sprintf("--runtime-class-names=%s", "contrast-cc"),
-		fmt.Sprintf("--rego-rules-path=%s", regoPath),
-		fmt.Sprintf("--json-settings-path=%s", policyPath),
-		fmt.Sprintf("--yaml-file=%s", yamlPath),
-		fmt.Sprintf("--layers-cache-file-path=%s", genpolicyCachePath),
-	}
-	genpolicy := exec.CommandContext(ctx, genpolicyPath, args...)
-	genpolicy.Env = append(genpolicy.Env, "RUST_LOG=info", "RUST_BACKTRACE=1")
-
-	logFilter := newLogTranslator(logger)
-	defer logFilter.stop()
-	var stdout bytes.Buffer
-	genpolicy.Stdout = &stdout
-	genpolicy.Stderr = logFilter
-
-	if err := genpolicy.Run(); err != nil {
-		var exitErr *exec.ExitError
-		if errors.As(err, &exitErr) {
-			return [32]byte{}, fmt.Errorf("genpolicy failed with exit code %d", exitErr.ExitCode())
-		}
-		return [32]byte{}, fmt.Errorf("genpolicy failed: %w", err)
-	}
-
-	if stdout.Len() == 0 {
-		logger.Info("Policy output is empty, ignoring the file", "yamlPath", yamlPath)
-		return [32]byte{}, nil
-	}
-	policyHash := sha256.Sum256(stdout.Bytes())
-
-	return policyHash, nil
-}
-
 func generateWorkloadOwnerKey(flags *generateFlags) error {
 	if flags.disableUpdates || len(flags.workloadOwnerKeys) != 1 {
 		// No need to generate keys
@@ -657,11 +557,3 @@ func createFileWithDefault(path string, perm fs.FileMode, dflt func() ([]byte, e
 	_, err = file.Write(content)
 	return err
 }
-
-func installDir() (string, error) {
-	home, err := os.UserHomeDir()
-	if err != nil {
-		return "", err
-	}
-	return filepath.Join(home, ".contrast"), nil
-}

cli/genpolicy/assets/genpolicy-rules.rego

@@ -0,0 +1 @@
+# THIS FILE IS REPLACED DURING BUILD AND ONLY HERE TO SATISFY GO TOOLING

cli/genpolicy/config.go

@@ -0,0 +1,44 @@
+// Copyright 2024 Edgeless Systems GmbH
+// SPDX-License-Identifier: AGPL-3.0-only
+
+package genpolicy
+
+import (
+	_ "embed"
+
+	"github.com/edgelesssys/contrast/internal/platforms"
+)
+
+var (
+	//go:embed assets/genpolicy
+	genpolicyBin []byte
+	//go:embed assets/genpolicy-settings.json
+	defaultGenpolicySettings []byte
+	//go:embed assets/genpolicy-rules.rego
+	aksCloudHypervisorSNPRules []byte
+	//go:embed assets/allow-all.rego
+	permissiveRules []byte
+)
+
+// Config contains configuration files for genpolicy.
+type Config struct {
+	// Rules is a Rego module that verifies agent requests.
+	Rules []byte
+	// Settings is a json config file that holds platform-specific configuration.
+	Settings []byte
+}
+
+// NewConfig selects the appropriate genpolicy configuration for the target platform.
+func NewConfig(platform platforms.Platform) *Config {
+	cfg := &Config{
+		Settings: defaultGenpolicySettings,
+	}
+	switch platform {
+	case platforms.AKSCloudHypervisorSNP:
+		cfg.Rules = aksCloudHypervisorSNPRules
+	default:
+		// TODO(burgerdev): use real rules for supported platforms.
+		cfg.Rules = permissiveRules
+	}
+	return cfg
+}

cli/genpolicy/genpolicy.go

@@ -0,0 +1,83 @@
+// Copyright 2024 Edgeless Systems GmbH
+// SPDX-License-Identifier: AGPL-3.0-only
+
+package genpolicy
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log/slog"
+	"os"
+	"os/exec"
+	"path/filepath"
+
+	"github.com/edgelesssys/contrast/internal/embedbin"
+)
+
+// Runner is a wrapper around the genpolicy tool.
+//
+// Create an instance with New(), call Run() to execute the tool, and call
+// Teardown() afterwards to clean up temporary files.
+type Runner struct {
+	genpolicy embedbin.Installed
+
+	rulesPath    string
+	settingsPath string
+	cachePath    string
+}
+
+// New creates a new Runner for the given configuration.
+func New(rulesPath, settingsPath, cachePath string) (*Runner, error) {
+	e := embedbin.New()
+	genpolicy, err := e.Install("", genpolicyBin)
+	if err != nil {
+		return nil, fmt.Errorf("installing genpolicy: %w", err)
+	}
+	if err := os.MkdirAll(filepath.Dir(cachePath), os.ModePerm); err != nil {
+		return nil, fmt.Errorf("creating cache file: %w", err)
+	}
+
+	runner := &Runner{
+		genpolicy:    genpolicy,
+		rulesPath:    rulesPath,
+		settingsPath: settingsPath,
+		cachePath:    cachePath,
+	}
+
+	return runner, nil
+}
+
+// Run runs the tool on the given yaml.
+//
+// Run can be called more than once.
+func (r *Runner) Run(ctx context.Context, yamlPath string, logger *slog.Logger) error {
+	args := []string{
+		"--runtime-class-names=contrast-cc",
+		"--rego-rules-path=" + r.rulesPath,
+		"--json-settings-path=" + r.settingsPath,
+		"--layers-cache-file-path=" + r.cachePath,
+		"--yaml-file=" + yamlPath,
+	}
+	genpolicy := exec.CommandContext(ctx, r.genpolicy.Path(), args...)
+	genpolicy.Env = append(genpolicy.Env, "RUST_LOG=info", "RUST_BACKTRACE=1")
+
+	logFilter := newLogTranslator(logger)
+	defer logFilter.stop()
+	genpolicy.Stdout = io.Discard
+	genpolicy.Stderr = logFilter
+
+	if err := genpolicy.Run(); err != nil {
+		return fmt.Errorf("running genpolicy: %w", err)
+	}
+
+	return nil
+}
+
+// Teardown cleans up temporary files and should be called after the last Run.
+func (r *Runner) Teardown() error {
+	if r.genpolicy != nil {
+		return r.genpolicy.Uninstall()
+	}
+	return nil
+}