From e3b371881f293a78d073cb865f5c253d5eda00ec Mon Sep 17 00:00:00 2001
From: Markus Rudy <mr@edgeless.systems>
Date: Wed, 20 Mar 2024 16:41:02 +0100
Subject: [PATCH] cli: don't allow empty coordinator policy hash

---
 cli/cmd/set.go    | 24 ++++++++++++++++++------
 cli/cmd/verify.go |  9 ++-------
 go.mod            |  2 +-
 3 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/cli/cmd/set.go b/cli/cmd/set.go
index 16b9ac64ac..1f0aacf2af 100644
--- a/cli/cmd/set.go
+++ b/cli/cmd/set.go
@@ -25,6 +25,7 @@ import (
 	"github.com/edgelesssys/contrast/internal/spinner"
 	"github.com/edgelesssys/contrast/internal/userapi"
 	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
@@ -52,7 +53,7 @@ func NewSetCmd() *cobra.Command {
 	cmd.Flags().StringP("manifest", "m", manifestFilename, "path to manifest (.json) file")
 	cmd.Flags().StringP("coordinator", "c", "", "endpoint the coordinator can be reached at")
 	must(cobra.MarkFlagRequired(cmd.Flags(), "coordinator"))
-	cmd.Flags().String("coordinator-policy-hash", DefaultCoordinatorPolicyHash, "expected policy hash of the coordinator, will not be checked if empty")
+	cmd.Flags().String("coordinator-policy-hash", DefaultCoordinatorPolicyHash, "override the expected policy hash of the coordinator")
 	cmd.Flags().String("workload-owner-key", workloadOwnerPEM, "path to workload owner key (.pem) file")
 
 	return cmd
@@ -157,6 +158,21 @@ type setFlags struct {
 	workspaceDir         string
 }
 
+func decodeCoordinatorPolicyHash(flags *pflag.FlagSet) ([]byte, error) {
+	hexEncoded, err := flags.GetString("coordinator-policy-hash")
+	if err != nil {
+		return nil, fmt.Errorf("getting coordinator-policy-hash flag: %w", err)
+	}
+	hash, err := hex.DecodeString(hexEncoded)
+	if err != nil {
+		return nil, fmt.Errorf("hex-decoding coordinator-policy-hash flag: %w", err)
+	}
+	if len(hash) != 32 {
+		return nil, fmt.Errorf("coordinator-policy-hash must be exactly 32 hex-encoded bytes, got %d", len(hash))
+	}
+	return hash, nil
+}
+
 func parseSetFlags(cmd *cobra.Command) (*setFlags, error) {
 	flags := &setFlags{}
 	var err error
@@ -169,14 +185,10 @@ func parseSetFlags(cmd *cobra.Command) (*setFlags, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to get coordinator flag: %w", err)
 	}
-	policyString, err := cmd.Flags().GetString("coordinator-policy-hash")
+	flags.policy, err = decodeCoordinatorPolicyHash(cmd.Flags())
 	if err != nil {
 		return nil, fmt.Errorf("failed to get coordinator-policy-hash flag: %w", err)
 	}
-	flags.policy, err = hex.DecodeString(policyString)
-	if err != nil {
-		return nil, fmt.Errorf("hex-decoding coordinator-policy-hash flag: %w", err)
-	}
 	flags.workloadOwnerKeyPath, err = cmd.Flags().GetString("workload-owner-key")
 	if err != nil {
 		return nil, fmt.Errorf("getting workload-owner-key flag: %w", err)
diff --git a/cli/cmd/verify.go b/cli/cmd/verify.go
index 8b9c9b8ee8..96cb0ef641 100644
--- a/cli/cmd/verify.go
+++ b/cli/cmd/verify.go
@@ -2,7 +2,6 @@ package cmd
 
 import (
 	"crypto/sha256"
-	"encoding/hex"
 	"fmt"
 	"net"
 	"os"
@@ -43,7 +42,7 @@ func NewVerifyCmd() *cobra.Command {
 	cmd.Flags().String("workspace-dir", verifyDir, "directory to write files to, if not set explicitly to another location")
 	cmd.Flags().StringP("coordinator", "c", "", "endpoint the coordinator can be reached at")
 	must(cobra.MarkFlagRequired(cmd.Flags(), "coordinator"))
-	cmd.Flags().String("coordinator-policy-hash", DefaultCoordinatorPolicyHash, "expected policy hash of the coordinator, will not be checked if empty")
+	cmd.Flags().String("coordinator-policy-hash", DefaultCoordinatorPolicyHash, "override the expected policy hash of the coordinator")
 
 	return cmd
 }
@@ -126,14 +125,10 @@ func parseVerifyFlags(cmd *cobra.Command) (*verifyFlags, error) {
 	if err != nil {
 		return nil, err
 	}
-	policyString, err := cmd.Flags().GetString("coordinator-policy-hash")
+	policy, err := decodeCoordinatorPolicyHash(cmd.Flags())
 	if err != nil {
 		return nil, err
 	}
-	policy, err := hex.DecodeString(policyString)
-	if err != nil {
-		return nil, fmt.Errorf("hex-decoding coordinator-policy-hash flag: %w", err)
-	}
 
 	return &verifyFlags{
 		coordinator:  coordinator,
diff --git a/go.mod b/go.mod
index 5e366d3746..538c68efae 100644
--- a/go.mod
+++ b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/google/go-sev-guest v0.11.0
 	github.com/spf13/afero v1.11.0
 	github.com/spf13/cobra v1.8.0
+	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.9.0
 	go.uber.org/goleak v1.3.0
 	golang.org/x/exp v0.0.0-20240222234643-814bf88cf225
@@ -48,7 +49,6 @@ require (
 	github.com/pborman/uuid v1.2.1 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/spf13/pflag v1.0.5 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
 	golang.org/x/crypto v0.19.0 // indirect
 	golang.org/x/net v0.21.0 // indirect