From d91b12b561016ffd93d803d58e7093b81bd743f7 Mon Sep 17 00:00:00 2001
From: Malte Poll <1780588+malt3@users.noreply.github.com>
Date: Thu, 8 Feb 2024 09:48:24 +0100
Subject: [PATCH] cli: use workload owner key during set command

---
 cli/set.go | 72 +++++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 68 insertions(+), 4 deletions(-)

diff --git a/cli/set.go b/cli/set.go
index 3997c461c2..c17f25bff1 100644
--- a/cli/set.go
+++ b/cli/set.go
@@ -2,10 +2,15 @@ package main
 
 import (
 	"context"
+	"crypto/ecdsa"
+	"crypto/sha256"
+	"crypto/x509"
 	"encoding/hex"
 	"encoding/json"
+	"encoding/pem"
 	"fmt"
 	"io"
+	"log/slog"
 	"net"
 	"os"
 	"time"
@@ -45,6 +50,7 @@ func newSetCmd() *cobra.Command {
 	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("workload-owner-key", workloadOwnerPEM, "path to workload owner key (.pem) file")
 
 	return cmd
 }
@@ -69,6 +75,11 @@ func runSet(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to unmarshal manifest: %w", err)
 	}
 
+	workloadOwnerKey, err := loadWorkloadOwnerKey(flags.workloadOwnerKeyPath, m, log)
+	if err != nil {
+		return fmt.Errorf("loading workload owner key: %w", err)
+	}
+
 	paths, err := findGenerateTargets(args, log)
 	if err != nil {
 		return fmt.Errorf("finding yaml files: %w", err)
@@ -92,7 +103,7 @@ func runSet(cmd *cobra.Command, args []string) error {
 	kdsCache := fsstore.New(kdsDir, log.WithGroup("kds-cache"))
 	kdsGetter := snp.NewCachedHTTPSGetter(kdsCache, snp.NeverGCTicker, log.WithGroup("kds-getter"))
 	validator := snp.NewValidator(validateOptsGen, kdsGetter, log.WithGroup("snp-validator"))
-	dialer := dialer.New(atls.NoIssuer, validator, &net.Dialer{})
+	dialer := dialer.NewWithKey(atls.NoIssuer, validator, &net.Dialer{}, workloadOwnerKey)
 
 	conn, err := dialer.Dial(cmd.Context(), flags.coordinator)
 	if err != nil {
@@ -124,9 +135,10 @@ func runSet(cmd *cobra.Command, args []string) error {
 }
 
 type setFlags struct {
-	manifestPath string
-	coordinator  string
-	policy       []byte
+	manifestPath         string
+	coordinator          string
+	policy               []byte
+	workloadOwnerKeyPath string
 }
 
 func parseSetFlags(cmd *cobra.Command) (*setFlags, error) {
@@ -149,6 +161,10 @@ func parseSetFlags(cmd *cobra.Command) (*setFlags, error) {
 	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)
+	}
 
 	return flags, nil
 }
@@ -161,6 +177,54 @@ func policyMapToBytesList(m map[string]deployment) [][]byte {
 	return policies
 }
 
+func loadWorkloadOwnerKey(path string, manifst manifest.Manifest, log *slog.Logger) (*ecdsa.PrivateKey, error) {
+	key, err := os.ReadFile(path)
+	if os.IsNotExist(err) {
+		if len(manifst.WorkloadOwnerKeys) == 0 {
+			log.Warn("No workload owner keys in manifest nor private key found. Further manifest updates will be rejected by the coordinator.")
+		} else {
+			log.Warn("No workload owner private key found. Setting the manifest may fail.")
+		}
+		return nil, nil
+	}
+	if err != nil {
+		return nil, fmt.Errorf("reading workload owner key: %w", err)
+	}
+	pemBlock, _ := pem.Decode(key)
+	if pemBlock == nil {
+		return nil, fmt.Errorf("decoding workload owner key: %w", err)
+	}
+	if pemBlock.Type != "EC PRIVATE KEY" {
+		return nil, fmt.Errorf("workload owner key is not an EC private key")
+	}
+	workloadOwnerKey, err := x509.ParseECPrivateKey(pemBlock.Bytes)
+	if err != nil {
+		return nil, fmt.Errorf("parsing workload owner key: %w", err)
+	}
+	pubKey, err := x509.MarshalPKIXPublicKey(&workloadOwnerKey.PublicKey)
+	if err != nil {
+		return nil, fmt.Errorf("marshaling public key: %w", err)
+	}
+	ownerKeyHash := sha256.Sum256(pubKey)
+	ownerKeyHex := manifest.NewHexString(ownerKeyHash[:])
+	if len(manifst.WorkloadOwnerKeys) == 0 {
+		log.Warn("No workload owner keys in manifest. Further manifest updates will be rejected by the coordinator")
+		return workloadOwnerKey, nil
+	}
+	log.Debug("Workload owner keys in manifest", "keys", manifst.WorkloadOwnerKeys)
+	var found bool
+	for _, k := range manifst.WorkloadOwnerKeys {
+		if k == ownerKeyHex {
+			found = true
+			break
+		}
+	}
+	if !found {
+		log.Warn("Workload owner key not found in manifest. This may lock you out from further updates")
+	}
+	return workloadOwnerKey, nil
+}
+
 func setLoop(
 	ctx context.Context, client coordapi.CoordAPIClient, out io.Writer, req *coordapi.SetManifestRequest,
 ) (resp *coordapi.SetManifestResponse, retErr error) {