ipn/store: automatically migrate between plaintext and encrypted state (#16318)

Add a new `--encrypt-state` flag to `cmd/tailscaled`. Based on that
flag, migrate the existing state file to/from encrypted format if
needed.

Updates #15830

Signed-off-by: Andrew Lytvynov <awly@tailscale.com>

ipn/store/awsstore/store_aws.go

@@ -10,6 +10,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"iter"
 	"net/url"
 	"regexp"
 	"strings"
@@ -253,3 +254,7 @@ func (s *awsStore) persistState() error {
 	_, err = s.ssmClient.PutParameter(context.TODO(), in)
 	return err
 }
+
+func (s *awsStore) All() iter.Seq2[ipn.StateKey, []byte] {
+	return s.memory.All()
+}

ipn/store/kubestore/store_kube.go

@@ -7,6 +7,7 @@ package kubestore
 import (
 	"context"
 	"fmt"
+	"iter"
 	"log"
 	"net"
 	"os"
@@ -428,3 +429,7 @@ func sanitizeKey[T ~string](k T) string {
 		return '_'
 	}, string(k))
 }
+
+func (s *Store) All() iter.Seq2[ipn.StateKey, []byte] {
+	return s.memory.All()
+}

ipn/store/mem/store_mem.go

@@ -7,6 +7,7 @@ package mem
 import (
 	"bytes"
 	"encoding/json"
+	"iter"
 	"sync"
 
 	xmaps "golang.org/x/exp/maps"
@@ -85,3 +86,16 @@ func (s *Store) ExportToJSON() ([]byte, error) {
 	}
 	return json.MarshalIndent(s.cache, "", "  ")
 }
+
+func (s *Store) All() iter.Seq2[ipn.StateKey, []byte] {
+	return func(yield func(ipn.StateKey, []byte) bool) {
+		s.mu.Lock()
+		defer s.mu.Unlock()
+
+		for k, v := range s.cache {
+			if !yield(k, v) {
+				break
+			}
+		}
+	}
+}

ipn/store/stores.go

@@ -7,10 +7,14 @@ package store
 import (
 	"bytes"
 	"encoding/json"
+	"errors"
 	"fmt"
+	"iter"
+	"maps"
 	"os"
 	"path/filepath"
 	"runtime"
+	"slices"
 	"strings"
 	"sync"
 
@@ -20,6 +24,7 @@ import (
 	"tailscale.com/paths"
 	"tailscale.com/types/logger"
 	"tailscale.com/util/mak"
+	"tailscale.com/util/testenv"
 )
 
 // Provider returns a StateStore for the provided path.
@@ -32,6 +37,9 @@ func init() {
 
 var knownStores map[string]Provider
 
+// TPMPrefix is the path prefix used for TPM-encrypted StateStore.
+const TPMPrefix = "tpmseal:"
+
 // New returns a StateStore based on the provided arg
 // and registered stores.
 // The arg is of the form "prefix:rest", where prefix was previously
@@ -53,12 +61,23 @@ func New(logf logger.Logf, path string) (ipn.StateStore, error) {
 		if strings.HasPrefix(path, prefix) {
 			// We can't strip the prefix here as some NewStoreFunc (like arn:)
 			// expect the prefix.
+			if prefix == TPMPrefix {
+				if runtime.GOOS == "windows" {
+					path = TPMPrefix + TryWindowsAppDataMigration(logf, strings.TrimPrefix(path, TPMPrefix))
+				}
+				if err := maybeMigrateLocalStateFile(logf, path); err != nil {
+					return nil, fmt.Errorf("failed to migrate existing state file to TPM-sealed format: %w", err)
+				}
+			}
 			return sf(logf, path)
 		}
 	}
 	if runtime.GOOS == "windows" {
 		path = TryWindowsAppDataMigration(logf, path)
 	}
+	if err := maybeMigrateLocalStateFile(logf, path); err != nil {
+		return nil, fmt.Errorf("failed to migrate existing TPM-sealed state file to plaintext format: %w", err)
+	}
 	return NewFileStore(logf, path)
 }
 
@@ -77,6 +96,29 @@ func Register(prefix string, fn Provider) {
 	mak.Set(&knownStores, prefix, fn)
 }
 
+// RegisterForTest registers a prefix to be used for NewStore in tests. An
+// existing registered prefix will be replaced.
+func RegisterForTest(t testenv.TB, prefix string, fn Provider) {
+	if len(prefix) == 0 {
+		panic("prefix is empty")
+	}
+	old := maps.Clone(knownStores)
+	t.Cleanup(func() { knownStores = old })
+
+	mak.Set(&knownStores, prefix, fn)
+}
+
+// HasKnownProviderPrefix reports whether path uses one of the registered
+// Provider prefixes.
+func HasKnownProviderPrefix(path string) bool {
+	for prefix := range knownStores {
+		if strings.HasPrefix(path, prefix) {
+			return true
+		}
+	}
+	return false
+}
+
 // TryWindowsAppDataMigration attempts to copy the Windows state file
 // from its old location to the new location. (Issue 2856)
 //
@@ -179,3 +221,101 @@ func (s *FileStore) WriteState(id ipn.StateKey, bs []byte) error {
 	}
 	return atomicfile.WriteFile(s.path, bs, 0600)
 }
+
+func (s *FileStore) All() iter.Seq2[ipn.StateKey, []byte] {
+	return func(yield func(ipn.StateKey, []byte) bool) {
+		s.mu.Lock()
+		defer s.mu.Unlock()
+
+		for k, v := range s.cache {
+			if !yield(k, v) {
+				break
+			}
+		}
+	}
+}
+
+func maybeMigrateLocalStateFile(logf logger.Logf, path string) error {
+	path, toTPM := strings.CutPrefix(path, TPMPrefix)
+
+	// Extract JSON keys from the file on disk and guess what kind it is.
+	bs, err := os.ReadFile(path)
+	if err != nil {
+		if errors.Is(err, os.ErrNotExist) {
+			return nil
+		}
+		return err
+	}
+	var content map[string]any
+	if err := json.Unmarshal(bs, &content); err != nil {
+		return fmt.Errorf("failed to unmarshal %q: %w", path, err)
+	}
+	keys := slices.Sorted(maps.Keys(content))
+	tpmKeys := []string{"key", "nonce", "data"}
+	slices.Sort(tpmKeys)
+	// TPM-sealed files will have exactly these keys.
+	existingFileSealed := slices.Equal(keys, tpmKeys)
+	// Plaintext files for nodes that registered at least once will have this
+	// key, plus other dynamic ones.
+	_, existingFilePlaintext := content["_machinekey"]
+	isTPM := existingFileSealed && !existingFilePlaintext
+
+	if isTPM == toTPM {
+		// No migration needed.
+		return nil
+	}
+
+	newTPMStore, ok := knownStores[TPMPrefix]
+	if !ok {
+		return errors.New("this build does not support TPM integration")
+	}
+
+	// Open from (old format) and to (new format) stores for migration. The
+	// "to" store will be at tmpPath.
+	var from, to ipn.StateStore
+	tmpPath := path + ".tmp"
+	if toTPM {
+		// Migrate plaintext file to be TPM-sealed.
+		from, err = NewFileStore(logf, path)
+		if err != nil {
+			return fmt.Errorf("NewFileStore(%q): %w", path, err)
+		}
+		to, err = newTPMStore(logf, TPMPrefix+tmpPath)
+		if err != nil {
+			return fmt.Errorf("newTPMStore(%q): %w", tmpPath, err)
+		}
+	} else {
+		// Migrate TPM-selaed file to plaintext.
+		from, err = newTPMStore(logf, TPMPrefix+path)
+		if err != nil {
+			return fmt.Errorf("newTPMStore(%q): %w", path, err)
+		}
+		to, err = NewFileStore(logf, tmpPath)
+		if err != nil {
+			return fmt.Errorf("NewFileStore(%q): %w", tmpPath, err)
+		}
+	}
+	defer os.Remove(tmpPath)
+
+	// Copy all the items. This is pretty inefficient, because both stores
+	// write the file to disk for each WriteState, but that's ok for a one-time
+	// migration.
+	for k, v := range from.All() {
+		if err := to.WriteState(k, v); err != nil {
+			return err
+		}
+	}
+
+	// Finally, overwrite the state file with the new one we created at
+	// tmpPath.
+	if err := atomicfile.Rename(tmpPath, path); err != nil {
+		return err
+	}
+
+	if toTPM {
+		logf("migrated %q from plaintext to TPM-sealed format", path)
+	} else {
+		logf("migrated %q from TPM-sealed to plaintext format", path)
+	}
+	return nil
+}