ipn/store: make StateStore.All optional (#16409)

This method is only needed to migrate between store.FileStore and
tpm.tpmStore. We can make a runtime type assertion instead of
implementing an unused method for every platform.

Updates #15830

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

cmd/tsconnect/src/lib/js-state-store.ts

@@ -10,7 +10,4 @@ export const sessionStateStorage: IPNStateStorage = {
   getState(id) {
     return window.sessionStorage[`ipn-state-${id}`] || ""
   },
-  all() {
-    return JSON.stringify(window.sessionStorage)
-  },
 }

cmd/tsconnect/src/types/wasm_js.d.ts

@@ -44,7 +44,6 @@ declare global {
   interface IPNStateStorage {
     setState(id: string, value: string): void
     getState(id: string): string
-    all(): string
   }
 
   type IPNConfig = {

cmd/tsconnect/wasm/wasm_js.go

@@ -15,7 +15,6 @@ import (
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
-	"iter"
 	"log"
 	"math/rand/v2"
 	"net"
@@ -580,29 +579,6 @@ func (s *jsStateStore) WriteState(id ipn.StateKey, bs []byte) error {
 	return nil
 }
 
-func (s *jsStateStore) All() iter.Seq2[ipn.StateKey, []byte] {
-	return func(yield func(ipn.StateKey, []byte) bool) {
-		jsValue := s.jsStateStorage.Call("all")
-		if jsValue.String() == "" {
-			return
-		}
-		buf, err := hex.DecodeString(jsValue.String())
-		if err != nil {
-			return
-		}
-		var state map[string][]byte
-		if err := json.Unmarshal(buf, &state); err != nil {
-			return
-		}
-
-		for k, v := range state {
-			if !yield(ipn.StateKey(k), v) {
-				break
-			}
-		}
-	}
-}
-
 func mapSlice[T any, M any](a []T, f func(T) M) []M {
 	n := make([]M, len(a))
 	for i, e := range a {

feature/tpm/tpm.go

@@ -217,6 +217,10 @@ func (s *tpmStore) All() iter.Seq2[ipn.StateKey, []byte] {
 	}
 }
 
+// Ensure tpmStore implements store.ExportableStore for migration to/from
+// store.FileStore.
+var _ store.ExportableStore = (*tpmStore)(nil)
+
 // The nested levels of encoding and encryption are confusing, so here's what's
 // going on in plain English.
 //

feature/tpm/tpm_test.go

@@ -9,6 +9,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"iter"
 	"maps"
 	"os"
 	"path/filepath"
@@ -20,8 +21,8 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"tailscale.com/ipn"
 	"tailscale.com/ipn/store"
-	"tailscale.com/ipn/store/mem"
 	"tailscale.com/types/logger"
+	"tailscale.com/util/mak"
 )
 
 func TestPropToString(t *testing.T) {
@@ -251,7 +252,9 @@ func TestMigrateStateToTPM(t *testing.T) {
 			if err != nil {
 				t.Fatalf("migration failed: %v", err)
 			}
-			gotContent := maps.Collect(s.All())
+			gotContent := maps.Collect(s.(interface {
+				All() iter.Seq2[ipn.StateKey, []byte]
+			}).All())
 			if diff := cmp.Diff(content, gotContent); diff != "" {
 				t.Errorf("unexpected content after migration, diff:\n%s", diff)
 			}
@@ -285,7 +288,7 @@ func tpmSupported() bool {
 
 type mockTPMSealProvider struct {
 	path string
-	mem.Store
+	data map[ipn.StateKey][]byte
 }
 
 func newMockTPMSeal(logf logger.Logf, path string) (ipn.StateStore, error) {
@@ -293,7 +296,7 @@ func newMockTPMSeal(logf logger.Logf, path string) (ipn.StateStore, error) {
 	if !ok {
 		return nil, fmt.Errorf("%q missing tpmseal: prefix", path)
 	}
-	s := &mockTPMSealProvider{path: path, Store: mem.Store{}}
+	s := &mockTPMSealProvider{path: path}
 	buf, err := os.ReadFile(path)
 	if errors.Is(err, os.ErrNotExist) {
 		return s, s.flushState()
@@ -312,24 +315,28 @@ func newMockTPMSeal(logf logger.Logf, path string) (ipn.StateStore, error) {
 	if data.Key == "" || data.Nonce == "" {
 		return nil, fmt.Errorf("%q missing key or nonce", path)
 	}
-	for k, v := range data.Data {
-		s.Store.WriteState(k, v)
-	}
+	s.data = data.Data
 	return s, nil
 }
 
-func (p *mockTPMSealProvider) WriteState(k ipn.StateKey, v []byte) error {
-	if err := p.Store.WriteState(k, v); err != nil {
-		return err
+func (p *mockTPMSealProvider) ReadState(k ipn.StateKey) ([]byte, error) {
+	return p.data[k], nil
 }
+
+func (p *mockTPMSealProvider) WriteState(k ipn.StateKey, v []byte) error {
+	mak.Set(&p.data, k, v)
 	return p.flushState()
 }
 
+func (p *mockTPMSealProvider) All() iter.Seq2[ipn.StateKey, []byte] {
+	return maps.All(p.data)
+}
+
 func (p *mockTPMSealProvider) flushState() error {
 	data := map[string]any{
 		"key":   "foo",
 		"nonce": "bar",
-		"data":  maps.Collect(p.Store.All()),
+		"data":  p.data,
 	}
 	buf, err := json.Marshal(data)
 	if err != nil {

ipn/store.go

@@ -8,7 +8,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"iter"
 	"net"
 	"strconv"
 )
@@ -84,11 +83,6 @@ type StateStore interface {
 	// instead, which only writes if the value is different from what's
 	// already in the store.
 	WriteState(id StateKey, bs []byte) error
-	// All returns an iterator over all StateStore keys. Using ReadState or
-	// WriteState is not safe while iterating and can lead to a deadlock.
-	// The order of keys in the iterator is not specified and may change
-	// between runs.
-	All() iter.Seq2[StateKey, []byte]
 }
 
 // WriteState is a wrapper around store.WriteState that only writes if

ipn/store/awsstore/store_aws.go

@@ -10,7 +10,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"iter"
 	"net/url"
 	"regexp"
 	"strings"
@@ -254,7 +253,3 @@ 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,7 +7,6 @@ package kubestore
 import (
 	"context"
 	"fmt"
-	"iter"
 	"log"
 	"net"
 	"os"
@@ -429,7 +428,3 @@ 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,7 +7,6 @@ package mem
 import (
 	"bytes"
 	"encoding/json"
-	"iter"
 	"sync"
 
 	xmaps "golang.org/x/exp/maps"
@@ -86,16 +85,3 @@ 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

@@ -235,6 +235,23 @@ func (s *FileStore) All() iter.Seq2[ipn.StateKey, []byte] {
 	}
 }
 
+// Ensure FileStore implements ExportableStore for migration to/from
+// tpm.tpmStore.
+var _ ExportableStore = (*FileStore)(nil)
+
+// ExportableStore is an ipn.StateStore that can export all of its contents.
+// This interface is optional to implement, and used for migrating the state
+// between different store implementations.
+type ExportableStore interface {
+	ipn.StateStore
+
+	// All returns an iterator over all store keys. Using ReadState or
+	// WriteState is not safe while iterating and can lead to a deadlock. The
+	// order of keys in the iterator is not specified and may change between
+	// runs.
+	All() iter.Seq2[ipn.StateKey, []byte]
+}
+
 func maybeMigrateLocalStateFile(logf logger.Logf, path string) error {
 	path, toTPM := strings.CutPrefix(path, TPMPrefix)
 
@@ -297,10 +314,15 @@ func maybeMigrateLocalStateFile(logf logger.Logf, path string) error {
 	}
 	defer os.Remove(tmpPath)
 
+	fromExp, ok := from.(ExportableStore)
+	if !ok {
+		return fmt.Errorf("%T does not implement the exportableStore interface", from)
+	}
+
 	// 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() {
+	for k, v := range fromExp.All() {
 		if err := to.WriteState(k, v); err != nil {
 			return err
 		}