feature/tpm: implement ipn.StateStore using TPM sealing (#16030)

Updates #15830

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

feature/tpm/tpm_test.go

@@ -3,7 +3,17 @@
 
 package tpm
 
-import "testing"
+import (
+	"bytes"
+	"crypto/rand"
+	"errors"
+	"path/filepath"
+	"strconv"
+	"testing"
+
+	"tailscale.com/ipn"
+	"tailscale.com/ipn/store"
+)
 
 func TestPropToString(t *testing.T) {
 	for prop, want := range map[uint32]string{
@@ -17,3 +27,156 @@ func TestPropToString(t *testing.T) {
 		}
 	}
 }
+
+func skipWithoutTPM(t testing.TB) {
+	tpm, err := open()
+	if err != nil {
+		t.Skip("TPM not available")
+	}
+	tpm.Close()
+}
+
+func TestSealUnseal(t *testing.T) {
+	skipWithoutTPM(t)
+
+	data := make([]byte, 100*1024)
+	rand.Read(data)
+	var key [32]byte
+	rand.Read(key[:])
+
+	sealed, err := seal(t.Logf, decryptedData{Key: key, Data: data})
+	if err != nil {
+		t.Fatalf("seal: %v", err)
+	}
+	if bytes.Contains(sealed.Data, data) {
+		t.Fatalf("sealed data %q contains original input %q", sealed.Data, data)
+	}
+
+	unsealed, err := unseal(t.Logf, *sealed)
+	if err != nil {
+		t.Fatalf("unseal: %v", err)
+	}
+	if !bytes.Equal(data, unsealed.Data) {
+		t.Errorf("got unsealed data: %q, want: %q", unsealed, data)
+	}
+	if key != unsealed.Key {
+		t.Errorf("got unsealed key: %q, want: %q", unsealed.Key, key)
+	}
+}
+
+func TestStore(t *testing.T) {
+	skipWithoutTPM(t)
+
+	path := storePrefix + filepath.Join(t.TempDir(), "state")
+	store, err := newStore(t.Logf, path)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	checkState := func(t *testing.T, store ipn.StateStore, k ipn.StateKey, want []byte) {
+		got, err := store.ReadState(k)
+		if err != nil {
+			t.Errorf("ReadState(%q): %v", k, err)
+		}
+		if !bytes.Equal(want, got) {
+			t.Errorf("ReadState(%q): got %q, want %q", k, got, want)
+		}
+	}
+
+	k1, k2 := ipn.StateKey("k1"), ipn.StateKey("k2")
+	v1, v2 := []byte("v1"), []byte("v2")
+
+	t.Run("read-non-existent-key", func(t *testing.T) {
+		_, err := store.ReadState(k1)
+		if !errors.Is(err, ipn.ErrStateNotExist) {
+			t.Errorf("ReadState succeeded, want %v", ipn.ErrStateNotExist)
+		}
+	})
+
+	t.Run("read-write-k1", func(t *testing.T) {
+		if err := store.WriteState(k1, v1); err != nil {
+			t.Errorf("WriteState(%q, %q): %v", k1, v1, err)
+		}
+		checkState(t, store, k1, v1)
+	})
+
+	t.Run("read-write-k2", func(t *testing.T) {
+		if err := store.WriteState(k2, v2); err != nil {
+			t.Errorf("WriteState(%q, %q): %v", k2, v2, err)
+		}
+		checkState(t, store, k2, v2)
+	})
+
+	t.Run("update-k2", func(t *testing.T) {
+		v2 = []byte("new v2")
+		if err := store.WriteState(k2, v2); err != nil {
+			t.Errorf("WriteState(%q, %q): %v", k2, v2, err)
+		}
+		checkState(t, store, k2, v2)
+	})
+
+	t.Run("reopen-store", func(t *testing.T) {
+		store, err := newStore(t.Logf, path)
+		if err != nil {
+			t.Fatal(err)
+		}
+		checkState(t, store, k1, v1)
+		checkState(t, store, k2, v2)
+	})
+}
+
+func BenchmarkStore(b *testing.B) {
+	skipWithoutTPM(b)
+	b.StopTimer()
+
+	stores := make(map[string]ipn.StateStore)
+	key := ipn.StateKey(b.Name())
+
+	// Set up tpmStore
+	tpmStore, err := newStore(b.Logf, filepath.Join(b.TempDir(), "tpm.store"))
+	if err != nil {
+		b.Fatal(err)
+	}
+	if err := tpmStore.WriteState(key, []byte("-1")); err != nil {
+		b.Fatal(err)
+	}
+	stores["tpmStore"] = tpmStore
+
+	// Set up FileStore
+	fileStore, err := store.NewFileStore(b.Logf, filepath.Join(b.TempDir(), "file.store"))
+	if err != nil {
+		b.Fatal(err)
+	}
+	if err := fileStore.WriteState(key, []byte("-1")); err != nil {
+		b.Fatal(err)
+	}
+	stores["fileStore"] = fileStore
+
+	b.StartTimer()
+
+	for name, store := range stores {
+		b.Run(name, func(b *testing.B) {
+			b.Run("write-noop", func(b *testing.B) {
+				for range b.N {
+					if err := store.WriteState(key, []byte("-1")); err != nil {
+						b.Fatal(err)
+					}
+				}
+			})
+			b.Run("write", func(b *testing.B) {
+				for i := range b.N {
+					if err := store.WriteState(key, []byte(strconv.Itoa(i))); err != nil {
+						b.Fatal(err)
+					}
+				}
+			})
+			b.Run("read", func(b *testing.B) {
+				for range b.N {
+					if _, err := store.ReadState(key); err != nil {
+						b.Fatal(err)
+					}
+				}
+			})
+		})
+	}
+}