ipn/store/awsstore: allow providing a KMS keyjaxxstorm/ssm_kms

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>

4 files changed, 157 insertions, 43 deletions

diff --git a/ipn/store/awsstore/store_aws.go b/ipn/store/awsstore/store_aws.go
index 0fb78d45a..40bbbf037 100644
--- a/ipn/store/awsstore/store_aws.go
+++ b/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
 }
diff --git a/ipn/store/awsstore/store_aws_stub.go b/ipn/store/awsstore/store_aws_stub.go
deleted file mode 100644
index 8d2156ce9..000000000
--- a/ipn/store/awsstore/store_aws_stub.go
+++ /dev/null
@@ -1,18 +0,0 @@
-// Copyright (c) Tailscale Inc & AUTHORS
-// SPDX-License-Identifier: BSD-3-Clause
-
-//go:build !linux || ts_omit_aws
-
-package awsstore
-
-import (
-	"fmt"
-	"runtime"
-
-	"tailscale.com/ipn"
-	"tailscale.com/types/logger"
-)
-
-func New(logger.Logf, string) (ipn.StateStore, error) {
-	return nil, fmt.Errorf("AWS store is not supported on %v", runtime.GOOS)
-}
diff --git a/ipn/store/awsstore/store_aws_test.go b/ipn/store/awsstore/store_aws_test.go
index f6c8fedb3..3382635a7 100644
--- a/ipn/store/awsstore/store_aws_test.go
+++ b/ipn/store/awsstore/store_aws_test.go
@@ -1,7 +1,7 @@
 // Copyright (c) Tailscale Inc & AUTHORS
 // SPDX-License-Identifier: BSD-3-Clause
 
-//go:build linux
+//go:build linux && !ts_omit_aws
 
 package awsstore
 
@@ -65,7 +65,11 @@ func TestNewAWSStore(t *testing.T) {
 		Resource:  "parameter/foo",
 	}
 
-	s, err := newStore(storeParameterARN.String(), mc)
+	opts := storeOptions{
+		kmsKey: "arn:aws:kms:eu-west-1:123456789:key/MyCustomKey",
+	}
+
+	s, err := newStore(storeParameterARN.String(), opts, mc)
 	if err != nil {
 		t.Fatalf("creating aws store failed: %v", err)
 	}
@@ -73,7 +77,7 @@ func TestNewAWSStore(t *testing.T) {
 
 	// Build a brand new file store and check that both IDs written
 	// above are still there.
-	s2, err := newStore(storeParameterARN.String(), mc)
+	s2, err := newStore(storeParameterARN.String(), opts, mc)
 	if err != nil {
 		t.Fatalf("creating second aws store failed: %v", err)
 	}
@@ -162,3 +166,54 @@ func testStoreSemantics(t *testing.T, store ipn.StateStore) {
 		}
 	}
 }
+
+func TestParseARNAndOpts(t *testing.T) {
+	tests := []struct {
+		name    string
+		arg     string
+		wantARN string
+		wantKey string
+	}{
+		{
+			name:    "no-key",
+			arg:     "arn:aws:ssm:us-east-1:123456789012:parameter/myTailscaleParam",
+			wantARN: "arn:aws:ssm:us-east-1:123456789012:parameter/myTailscaleParam",
+		},
+		{
+			name:    "custom-key",
+			arg:     "arn:aws:ssm:us-east-1:123456789012:parameter/myTailscaleParam?kmsKey=alias/MyCustomKey",
+			wantARN: "arn:aws:ssm:us-east-1:123456789012:parameter/myTailscaleParam",
+			wantKey: "alias/MyCustomKey",
+		},
+		{
+			name:    "bare-name",
+			arg:     "arn:aws:ssm:us-east-1:123456789012:parameter/myTailscaleParam?kmsKey=Bare",
+			wantARN: "arn:aws:ssm:us-east-1:123456789012:parameter/myTailscaleParam",
+			wantKey: "alias/Bare",
+		},
+		{
+			name:    "arn-arg",
+			arg:     "arn:aws:ssm:us-east-1:123456789012:parameter/myTailscaleParam?kmsKey=arn:foo",
+			wantARN: "arn:aws:ssm:us-east-1:123456789012:parameter/myTailscaleParam",
+			wantKey: "arn:foo",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			arn, opts, err := ParseARNAndOpts(tt.arg)
+			if err != nil {
+				t.Fatalf("New: %v", err)
+			}
+			if arn != tt.wantARN {
+				t.Errorf("ARN = %q; want %q", arn, tt.wantARN)
+			}
+			var got storeOptions
+			for _, opt := range opts {
+				opt(&got)
+			}
+			if got.kmsKey != tt.wantKey {
+				t.Errorf("kmsKey = %q; want %q", got.kmsKey, tt.wantKey)
+			}
+		})
+	}
+}
diff --git a/ipn/store/store_aws.go b/ipn/store/store_aws.go
index e164f9de7..d39e84319 100644
--- a/ipn/store/store_aws.go
+++ b/ipn/store/store_aws.go
@@ -6,7 +6,9 @@
 package store
 
 import (
+	"tailscale.com/ipn"
 	"tailscale.com/ipn/store/awsstore"
+	"tailscale.com/types/logger"
 )
 
 func init() {
@@ -14,5 +16,11 @@ func init() {
 }
 
 func registerAWSStore() {
-	Register("arn:", awsstore.New)
+	Register("arn:", func(logf logger.Logf, arg string) (ipn.StateStore, error) {
+		ssmARN, opts, err := awsstore.ParseARNAndOpts(arg)
+		if err != nil {
+			return nil, err
+		}
+		return awsstore.New(logf, ssmARN, opts...)
+	})
 }