ipn/store/aws, cmd/tailscaled: add AWS SSM ipn.StateStore implementation

From https://github.com/tailscale/tailscale/pull/1919 with
edits by bradfitz@.

This change introduces a new storage provider for the state file. It
allows users to leverage AWS SSM parameter store natively within
tailscaled, like:

    $ tailscaled --state=arn:aws:ssm:eu-west-1:123456789:parameter/foo

Known limitations:
- it is not currently possible to specific a custom KMS key ID

RELNOTE=tailscaled on Linux supports using AWS SSM for state

Edits-By: Brad Fitzpatrick <bradfitz@tailscale.com>
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
Signed-off-by: Maxime VISONNEAU <maxime.visonneau@gmail.com>

ipn/ipnserver/server.go

@@ -35,6 +35,7 @@ import (
 	"tailscale.com/ipn"
 	"tailscale.com/ipn/ipnlocal"
 	"tailscale.com/ipn/localapi"
+	"tailscale.com/ipn/store/aws"
 	"tailscale.com/log/filelogger"
 	"tailscale.com/logtail/backoff"
 	"tailscale.com/net/netstat"
@@ -638,6 +639,7 @@ func Run(ctx context.Context, logf logger.Logf, logid string, getEngine func() (
 	var store ipn.StateStore
 	if opts.StatePath != "" {
 		const kubePrefix = "kube:"
+		const arnPrefix = "arn:"
 		path := opts.StatePath
 		switch {
 		case strings.HasPrefix(path, kubePrefix):
@@ -646,6 +648,11 @@ func Run(ctx context.Context, logf logger.Logf, logid string, getEngine func() (
 			if err != nil {
 				return fmt.Errorf("ipn.NewKubeStore(%q): %v", secretName, err)
 			}
+		case strings.HasPrefix(path, arnPrefix):
+			store, err = aws.NewStore(path)
+			if err != nil {
+				return fmt.Errorf("aws.NewStore(%q): %v", path, err)
+			}
 		default:
 			if runtime.GOOS == "windows" {
 				path = tryWindowsAppDataMigration(logf, path)

ipn/store.go

@@ -158,6 +158,26 @@ func (s *MemoryStore) WriteState(id StateKey, bs []byte) error {
 	return nil
 }
 
+// LoadFromJSON attempts to unmarshal json content into the
+// in-memory cache.
+func (s *MemoryStore) LoadFromJSON(data []byte) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return json.Unmarshal(data, &s.cache)
+}
+
+// ExportToJSON exports the content of the cache to
+// JSON formatted []byte.
+func (s *MemoryStore) ExportToJSON() ([]byte, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if len(s.cache) == 0 {
+		// Avoid "null" serialization.
+		return []byte("{}"), nil
+	}
+	return json.MarshalIndent(s.cache, "", "  ")
+}
+
 // FileStore is a StateStore that uses a JSON file for persistence.
 type FileStore struct {
 	path string

ipn/store/aws/store_aws.go

@@ -0,0 +1,175 @@
+// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build linux
+// +build linux
+
+// Package aws contains an AWS SSM StateStore implementation.
+package aws
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"regexp"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/aws/arn"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/ssm"
+	ssmTypes "github.com/aws/aws-sdk-go-v2/service/ssm/types"
+	"tailscale.com/ipn"
+)
+
+const (
+	parameterNameRxStr = `^parameter/(.*)`
+)
+
+var parameterNameRx = regexp.MustCompile(parameterNameRxStr)
+
+// awsSSMClient is an interface allowing us to mock the couple of
+// API calls we are leveraging with the AWSStore provider
+type awsSSMClient interface {
+	GetParameter(ctx context.Context,
+		params *ssm.GetParameterInput,
+		optFns ...func(*ssm.Options)) (*ssm.GetParameterOutput, error)
+
+	PutParameter(ctx context.Context,
+		params *ssm.PutParameterInput,
+		optFns ...func(*ssm.Options)) (*ssm.PutParameterOutput, error)
+}
+
+// store is a store which leverages AWS SSM parameter store
+// to persist the state
+type awsStore struct {
+	ssmClient awsSSMClient
+	ssmARN    arn.ARN
+
+	memory ipn.MemoryStore
+}
+
+// NewStore returns a new ipn.StateStore using the AWS SSM storage
+// location given by ssmARN.
+func NewStore(ssmARN string) (ipn.StateStore, error) {
+	return newStore(ssmARN, 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) {
+	s := &awsStore{
+		ssmClient: client,
+	}
+
+	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)
+	}
+
+	if s.ssmClient == nil {
+		var cfg aws.Config
+		if cfg, err = config.LoadDefaultConfig(
+			context.TODO(),
+			config.WithRegion(s.ssmARN.Region),
+		); err != nil {
+			return nil, err
+		}
+		s.ssmClient = ssm.NewFromConfig(cfg)
+	}
+
+	// Hydrate cache with the potentially current state
+	if err := s.LoadState(); err != nil {
+		return nil, err
+	}
+	return s, nil
+
+}
+
+// LoadState attempts to read the state from AWS SSM parameter store key.
+func (s *awsStore) LoadState() error {
+	param, err := s.ssmClient.GetParameter(
+		context.TODO(),
+		&ssm.GetParameterInput{
+			Name:           aws.String(s.ParameterName()),
+			WithDecryption: true,
+		},
+	)
+
+	if err != nil {
+		var pnf *ssmTypes.ParameterNotFound
+		if errors.As(err, &pnf) {
+			// Create the parameter as it does not exist yet
+			// and return directly as it is defacto empty
+			return s.persistState()
+		}
+		return err
+	}
+
+	// Load the content in-memory
+	return s.memory.LoadFromJSON([]byte(*param.Parameter.Value))
+}
+
+// ParameterName returns the parameter name extracted from
+// the provided ARN
+func (s *awsStore) ParameterName() (name string) {
+	values := parameterNameRx.FindStringSubmatch(s.ssmARN.Resource)
+	if len(values) == 2 {
+		name = values[1]
+	}
+	return
+}
+
+// String returns the awsStore and the ARN of the SSM parameter store
+// configured to store the state
+func (s *awsStore) String() string { return fmt.Sprintf("awsStore(%q)", s.ssmARN.String()) }
+
+// ReadState implements the Store interface.
+func (s *awsStore) ReadState(id ipn.StateKey) (bs []byte, err error) {
+	return s.memory.ReadState(id)
+}
+
+// WriteState implements the Store interface.
+func (s *awsStore) WriteState(id ipn.StateKey, bs []byte) (err error) {
+	// Write the state in-memory
+	if err = s.memory.WriteState(id, bs); err != nil {
+		return
+	}
+
+	// Persist the state in AWS SSM parameter store
+	return s.persistState()
+}
+
+// PersistState saves the states into the AWS SSM parameter store
+func (s *awsStore) persistState() error {
+	// Generate JSON from in-memory cache
+	bs, err := s.memory.ExportToJSON()
+	if err != nil {
+		return err
+	}
+
+	// Store in AWS SSM parameter store
+	_, err = s.ssmClient.PutParameter(
+		context.TODO(),
+		&ssm.PutParameterInput{
+			Name:      aws.String(s.ParameterName()),
+			Value:     aws.String(string(bs)),
+			Overwrite: true,
+			Tier:      ssmTypes.ParameterTierStandard,
+			Type:      ssmTypes.ParameterTypeSecureString,
+		},
+	)
+	return err
+}

ipn/store/aws/store_aws_stub.go

@@ -0,0 +1,19 @@
+// Copyright (c) 2021 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build !linux
+// +build !linux
+
+package aws
+
+import (
+	"fmt"
+	"runtime"
+
+	"tailscale.com/ipn"
+)
+
+func NewStore(string) (ipn.StateStore, error) {
+	return nil, fmt.Errorf("AWS store is not supported on %v", runtime.GOOS)
+}

ipn/store/aws/store_aws_test.go

@@ -0,0 +1,166 @@
+// Copyright (c) 2021 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build linux
+// +build linux
+
+package aws
+
+import (
+	"context"
+	"testing"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/aws/arn"
+	"github.com/aws/aws-sdk-go-v2/service/ssm"
+	ssmTypes "github.com/aws/aws-sdk-go-v2/service/ssm/types"
+	"tailscale.com/ipn"
+	"tailscale.com/tstest"
+)
+
+type mockedAWSSSMClient struct {
+	value string
+}
+
+func (sp *mockedAWSSSMClient) GetParameter(_ context.Context, input *ssm.GetParameterInput, _ ...func(*ssm.Options)) (*ssm.GetParameterOutput, error) {
+	output := new(ssm.GetParameterOutput)
+	if sp.value == "" {
+		return output, &ssmTypes.ParameterNotFound{}
+	}
+
+	output.Parameter = &ssmTypes.Parameter{
+		Value: aws.String(sp.value),
+	}
+
+	return output, nil
+}
+
+func (sp *mockedAWSSSMClient) PutParameter(_ context.Context, input *ssm.PutParameterInput, _ ...func(*ssm.Options)) (*ssm.PutParameterOutput, error) {
+	sp.value = *input.Value
+	return new(ssm.PutParameterOutput), nil
+}
+
+func TestAWSStoreString(t *testing.T) {
+	store := &awsStore{
+		ssmARN: arn.ARN{
+			Service:   "ssm",
+			Region:    "eu-west-1",
+			AccountID: "123456789",
+			Resource:  "parameter/foo",
+		},
+	}
+	want := "awsStore(\"arn::ssm:eu-west-1:123456789:parameter/foo\")"
+	if got := store.String(); got != want {
+		t.Errorf("AWSStore.String = %q; want %q", got, want)
+	}
+}
+
+func TestNewAWSStore(t *testing.T) {
+	tstest.PanicOnLog()
+
+	mc := &mockedAWSSSMClient{}
+	storeParameterARN := arn.ARN{
+		Service:   "ssm",
+		Region:    "eu-west-1",
+		AccountID: "123456789",
+		Resource:  "parameter/foo",
+	}
+
+	s, err := newStore(storeParameterARN.String(), mc)
+	if err != nil {
+		t.Fatalf("creating aws store failed: %v", err)
+	}
+	testStoreSemantics(t, s)
+
+	// Build a brand new file store and check that both IDs written
+	// above are still there.
+	s2, err := newStore(storeParameterARN.String(), mc)
+	if err != nil {
+		t.Fatalf("creating second aws store failed: %v", err)
+	}
+	store2 := s.(*awsStore)
+
+	// This is specific to the test, with the non-mocked API, LoadState() should
+	// have been already called and sucessful as no err is returned from NewAWSStore()
+	s2.(*awsStore).LoadState()
+
+	expected := map[ipn.StateKey]string{
+		"foo": "bar",
+		"baz": "quux",
+	}
+	for id, want := range expected {
+		bs, err := store2.ReadState(id)
+		if err != nil {
+			t.Errorf("reading %q (2nd store): %v", id, err)
+		}
+		if string(bs) != want {
+			t.Errorf("reading %q (2nd store): got %q, want %q", id, string(bs), want)
+		}
+	}
+}
+
+func testStoreSemantics(t *testing.T, store ipn.StateStore) {
+	t.Helper()
+
+	tests := []struct {
+		// if true, data is data to write. If false, data is expected
+		// output of read.
+		write bool
+		id    ipn.StateKey
+		data  string
+		// If write=false, true if we expect a not-exist error.
+		notExists bool
+	}{
+		{
+			id:        "foo",
+			notExists: true,
+		},
+		{
+			write: true,
+			id:    "foo",
+			data:  "bar",
+		},
+		{
+			id:   "foo",
+			data: "bar",
+		},
+		{
+			id:        "baz",
+			notExists: true,
+		},
+		{
+			write: true,
+			id:    "baz",
+			data:  "quux",
+		},
+		{
+			id:   "foo",
+			data: "bar",
+		},
+		{
+			id:   "baz",
+			data: "quux",
+		},
+	}
+
+	for _, test := range tests {
+		if test.write {
+			if err := store.WriteState(test.id, []byte(test.data)); err != nil {
+				t.Errorf("writing %q to %q: %v", test.data, test.id, err)
+			}
+		} else {
+			bs, err := store.ReadState(test.id)
+			if err != nil {
+				if test.notExists && err == ipn.ErrStateNotExist {
+					continue
+				}
+				t.Errorf("reading %q: %v", test.id, err)
+				continue
+			}
+			if string(bs) != test.data {
+				t.Errorf("reading %q: got %q, want %q", test.id, string(bs), test.data)
+			}
+		}
+	}
+}