ipn/store/awsstore: allow providing a KMS key

Implements a KMS input for AWS parameter to support encrypting Tailscale
state

Fixes #14765

Change-Id: I39c0fae4bfd60a9aec17c5ea6a61d0b57143d4ba
Co-authored-by: Brad Fitzpatrick <bradfitz@tailscale.com>
Signed-off-by: Lee Briggs <lee@leebriggs.co.uk>

ipn/store/awsstore/store_aws.go

@@ -10,7 +10,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net/url"
 	"regexp"
+	"strings"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/aws/arn"
@@ -28,6 +30,14 @@ const (
 
 var parameterNameRx = regexp.MustCompile(parameterNameRxStr)
 
+// Option defines a functional option type for configuring awsStore.
+type Option func(*storeOptions)
+
+// storeOptions holds optional settings for creating a new awsStore.
+type storeOptions struct {
+	kmsKey string
+}
+
 // awsSSMClient is an interface allowing us to mock the couple of
 // API calls we are leveraging with the AWSStore provider
 type awsSSMClient interface {
@@ -46,6 +56,10 @@ type awsStore struct {
 	ssmClient awsSSMClient
 	ssmARN    arn.ARN
 
+	// kmsKey is optional. If empty, the parameter is stored in plaintext.
+	// If non-empty, the parameter is encrypted with this KMS key.
+	kmsKey string
+
 	memory mem.Store
 }
 
@@ -57,30 +71,80 @@ type awsStore struct {
 // Tailscaled to only only store new state in-memory and
 // restarting Tailscaled can fail until you delete your state
 // from the AWS Parameter Store.
-func New(_ logger.Logf, ssmARN string) (ipn.StateStore, error) {
-	return newStore(ssmARN, nil)
+//
+// If you want to specify an optional KMS key,
+// pass one or more Option objects, e.g. awsstore.WithKeyID("alias/my-key").
+func New(_ logger.Logf, ssmARN string, opts ...Option) (ipn.StateStore, error) {
+	// Apply all options to an empty storeOptions
+	var so storeOptions
+	for _, opt := range opts {
+		opt(&so)
+	}
+
+	return newStore(ssmARN, so, nil)
+}
+
+// WithKeyID sets the KMS key to be used for encryption. It can be
+// a KeyID, an alias ("alias/my-key"), or a full ARN.
+//
+// If kmsKey is empty, the Option is a no-op.
+func WithKeyID(kmsKey string) Option {
+	return func(o *storeOptions) {
+		o.kmsKey = kmsKey
+	}
+}
+
+// ParseARNAndOpts parses an ARN and optional URL-encoded parameters
+// from arg.
+func ParseARNAndOpts(arg string) (ssmARN string, opts []Option, err error) {
+	ssmARN = arg
+
+	// Support optional ?url-encoded-parameters.
+	if s, q, ok := strings.Cut(arg, "?"); ok {
+		ssmARN = s
+		q, err := url.ParseQuery(q)
+		if err != nil {
+			return "", nil, err
+		}
+
+		for k := range q {
+			switch k {
+			default:
+				return "", nil, fmt.Errorf("unknown arn option parameter %q", k)
+			case "kmsKey":
+				// We allow an ARN, a key ID, or an alias name for kmsKeyID.
+				// If it doesn't look like an ARN and doesn't have a '/',
+				// prepend "alias/" for KMS alias references.
+				kmsKey := q.Get(k)
+				if kmsKey != "" &&
+					!strings.Contains(kmsKey, "/") &&
+					!strings.HasPrefix(kmsKey, "arn:") {
+					kmsKey = "alias/" + kmsKey
+				}
+				if kmsKey != "" {
+					opts = append(opts, WithKeyID(kmsKey))
+				}
+			}
+		}
+	}
+	return ssmARN, opts, nil
 }
 
 // newStore is NewStore, but for tests. If client is non-nil, it's
 // used instead of making one.
-func newStore(ssmARN string, client awsSSMClient) (ipn.StateStore, error) {
+func newStore(ssmARN string, so storeOptions, client awsSSMClient) (ipn.StateStore, error) {
 	s := &awsStore{
 		ssmClient: client,
+		kmsKey:    so.kmsKey,
 	}
 
 	var err error
-
-	// Parse the ARN
 	if s.ssmARN, err = arn.Parse(ssmARN); err != nil {
 		return nil, fmt.Errorf("unable to parse the ARN correctly: %v", err)
 	}
-
-	// Validate the ARN corresponds to the SSM service
 	if s.ssmARN.Service != "ssm" {
 		return nil, fmt.Errorf("invalid service %q, expected 'ssm'", s.ssmARN.Service)
 	}
-
-	// Validate the ARN corresponds to a parameter store resource
 	if !parameterNameRx.MatchString(s.ssmARN.Resource) {
 		return nil, fmt.Errorf("invalid resource %q, expected to match %v", s.ssmARN.Resource, parameterNameRxStr)
 	}
@@ -96,12 +160,11 @@ func newStore(ssmARN string, client awsSSMClient) (ipn.StateStore, error) {
 		s.ssmClient = ssm.NewFromConfig(cfg)
 	}
 
-	// Hydrate cache with the potentially current state
+	// Preload existing state, if any
 	if err := s.LoadState(); err != nil {
 		return nil, err
 	}
 	return s, nil
-
 }
 
 // LoadState attempts to read the state from AWS SSM parameter store key.
@@ -172,15 +235,21 @@ func (s *awsStore) persistState() error {
 	// which is free. However, if it exceeds 4kb it switches the parameter to advanced tiering
 	// doubling the capacity to 8kb per the following docs:
 	// https://aws.amazon.com/about-aws/whats-new/2019/08/aws-systems-manager-parameter-store-announces-intelligent-tiering-to-enable-automatic-parameter-tier-selection/
-	_, err = s.ssmClient.PutParameter(
-		context.TODO(),
-		&ssm.PutParameterInput{
-			Name:      aws.String(s.ParameterName()),
-			Value:     aws.String(string(bs)),
-			Overwrite: aws.Bool(true),
-			Tier:      ssmTypes.ParameterTierIntelligentTiering,
-			Type:      ssmTypes.ParameterTypeSecureString,
-		},
-	)
+	in := &ssm.PutParameterInput{
+		Name:      aws.String(s.ParameterName()),
+		Value:     aws.String(string(bs)),
+		Overwrite: aws.Bool(true),
+		Tier:      ssmTypes.ParameterTierIntelligentTiering,
+		Type:      ssmTypes.ParameterTypeSecureString,
+	}
+
+	// If kmsKey is specified, encrypt with that key
+	// NOTE: this input allows any alias, keyID or ARN
+	// If this isn't specified, AWS will use the default KMS key
+	if s.kmsKey != "" {
+		in.KeyId = aws.String(s.kmsKey)
+	}
+
+	_, err = s.ssmClient.PutParameter(context.TODO(), in)
 	return err
 }