state: undo a failed re-registration without reverting newer writes

hscontrol/state/auth_tagged_expiry_test.go

@@ -4,7 +4,9 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"net/netip"
 	"runtime"
+	"slices"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -2480,6 +2482,238 @@ func TestPersistNodeIncludesCurrentExpiry(t *testing.T) {
 	require.True(t, expiry.Equal(*stored.Expiry))
 }
 
+// TestPAKReregisterRollbackKeepsSameBatchWrites: admin writes queued behind a
+// failing re-registration share its NodeStore batch, so the rollback must
+// compare against what the registration wrote, not the node after the batch.
+func TestPAKReregisterRollbackKeepsSameBatchWrites(t *testing.T) {
+	s := newRetagTestState(t)
+	user := s.CreateUserForTest("rollback-user")
+	_, err := s.SetPolicy([]byte(`{"tagOwners":{
+    "tag:prod":["rollback-user@"],"tag:quarantine":["rollback-user@"]}}`))
+	require.NoError(t, err)
+	originalKey, err := s.CreatePreAuthKey(user.TypedID(), true, false, nil, nil)
+	require.NoError(t, err)
+	a := registerForPAKRereg(t, s, originalKey)
+	b := registerForPAKRereg(t, s, originalKey)
+
+	// Both rows are read unused; a spends the key, so b fails in the database.
+	sharedKey, err := s.CreatePreAuthKey(nil, false, false, nil, []string{"tag:prod"})
+	require.NoError(t, err)
+	pakA, err := s.GetPreAuthKeyByID(sharedKey.ID)
+	require.NoError(t, err)
+	pakB, err := s.GetPreAuthKeyByID(sharedKey.ID)
+	require.NoError(t, err)
+
+	reqA := a.regReq
+	reqA.Auth = &tailcfg.RegisterResponseAuth{AuthKey: sharedKey.Key}
+	_, err = a.reregister(a.node, pakA, reqA)
+	require.NoError(t, err)
+
+	reqB := b.regReq
+	reqB.Auth = &tailcfg.RegisterResponseAuth{AuthKey: sharedKey.Key}
+	reqB.NodeKey = key.NewNode().Public()
+
+	// Hold the writer so the registration and both admin writes queue behind it
+	// and land in one batch, in this order.
+	entered, release := make(chan struct{}), make(chan struct{})
+
+	releaseOnce := sync.OnceFunc(func() { close(release) })
+	defer releaseOnce()
+
+	go s.nodeStore.UpdateNode(a.node.ID(), func(*types.Node) {
+		close(entered)
+		<-release
+	})
+
+	<-entered
+
+	regDone := make(chan error, 1)
+
+	go func() {
+		_, err := b.reregister(b.node, pakB, reqB)
+		regDone <- err
+	}()
+
+	waitParkedOnWriteQueue(t, "reregisterNodeWithPAK")
+
+	tagsDone := make(chan error, 1)
+
+	go func() {
+		_, _, err := s.SetNodeTags(b.node.ID(), []string{"tag:quarantine"})
+		tagsDone <- err
+	}()
+
+	waitParkedOnWriteQueue(t, "SetNodeTags")
+
+	expiry := time.Now().Add(-time.Minute)
+	expiryDone := make(chan error, 1)
+
+	go func() {
+		_, _, err := s.SetNodeExpiry(b.node.ID(), &expiry)
+		expiryDone <- err
+	}()
+
+	waitParkedOnWriteQueue(t, "SetNodeExpiry")
+
+	releaseOnce()
+
+	require.ErrorIs(t, <-regDone, types.PAKError("authkey already used"))
+	require.NoError(t, <-tagsDone)
+	require.NoError(t, <-expiryDone)
+
+	ns, row := b.fields(t)
+	assert.Equal(t, []string{"tag:quarantine"}, ns.Tags, "admin tags must survive")
+	assert.Equal(t, expiry.UnixNano(), ns.Expiry, "admin expiry must survive")
+	assert.Equal(t, b.regReq.NodeKey, ns.NodeKey, "registration key must roll back")
+	assert.Equal(t, ns, row, "row must match NodeStore")
+}
+
+func TestPAKReregisterRollbackKeepsConcurrentWrites(t *testing.T) {
+	for _, update := range []string{"persisted_grant", "admin", "map"} {
+		t.Run(update, func(t *testing.T) {
+			dbPath, s, _ := persistTestSetup(t)
+			t.Cleanup(func() { _ = s.Close() })
+
+			user := s.CreateUserForTest("rollback-user")
+			_, err := s.SetPolicy([]byte(`{"tagOwners":{
+    "tag:prod":["rollback-user@"],"tag:quarantine":["rollback-user@"]}}`))
+			require.NoError(t, err)
+			originalKey, err := s.CreatePreAuthKey(user.TypedID(), true, false, nil, nil)
+			require.NoError(t, err)
+			a := registerForPAKRereg(t, s, originalKey)
+			b := registerForPAKRereg(t, s, originalKey)
+			sharedKey, err := s.CreatePreAuthKey(nil, false, false, nil, []string{"tag:prod"})
+			require.NoError(t, err)
+			pakA, err := s.GetPreAuthKeyByID(sharedKey.ID)
+			require.NoError(t, err)
+			pakB, err := s.GetPreAuthKeyByID(sharedKey.ID)
+			require.NoError(t, err)
+			require.False(t, pakA.Used)
+			require.False(t, pakB.Used)
+
+			reqA := a.regReq
+			reqA.Auth = &tailcfg.RegisterResponseAuth{AuthKey: sharedKey.Key}
+			_, err = a.reregister(a.node, pakA, reqA)
+			require.NoError(t, err)
+
+			reqB := b.regReq
+			reqB.Auth = &tailcfg.RegisterResponseAuth{AuthKey: sharedKey.Key}
+			reqB.NodeKey = key.NewNode().Public()
+			want := pakFieldsOf(b.node.AsStruct())
+			expiry := time.Now().Add(-time.Minute)
+			endpoint := netip.MustParseAddrPort("192.0.2.1:1234")
+
+			var pending []<-chan error
+
+			pool := s.db.DB.ConnPool
+			barrier := &beforeNodeTransaction{ConnPool: pool, before: func() {
+				published, ok := s.GetNodeByID(b.node.ID())
+				require.True(t, ok)
+				require.Equal(t, []string{"tag:prod"}, published.Tags().AsSlice())
+				// Another persist can write the published grant before registration
+				// takes persistMu. Reproduce that row at the transaction boundary.
+				require.NoError(t, s.db.DB.Select(nodeUpdateColumns).Updates(published.AsStruct()).Error)
+
+				switch update {
+				case "admin":
+					tagsDone := make(chan error, 1)
+					pending = append(pending, tagsDone)
+
+					go func() {
+						_, _, err := s.SetNodeTags(b.node.ID(), []string{"tag:quarantine"})
+						tagsDone <- err
+					}()
+
+					require.Eventually(t, func() bool {
+						n, ok := s.GetNodeByID(b.node.ID())
+						return ok && slices.Equal(n.Tags().AsSlice(), []string{"tag:quarantine"})
+					}, 5*time.Second, time.Millisecond)
+
+					expiryDone := make(chan error, 1)
+					pending = append(pending, expiryDone)
+
+					go func() {
+						_, _, err := s.SetNodeExpiry(b.node.ID(), &expiry)
+						expiryDone <- err
+					}()
+
+					require.Eventually(t, func() bool {
+						n, ok := s.GetNodeByID(b.node.ID())
+						return ok && n.IsExpired()
+					}, 5*time.Second, time.Millisecond)
+
+					want.Tags = []string{"tag:quarantine"}
+					want.UserID = nil
+					want.Expiry = expiry.UnixNano()
+				case "map":
+					mapDone := make(chan error, 1)
+					pending = append(pending, mapDone)
+
+					go func() {
+						_, err := s.UpdateNodeFromMapRequest(b.node.ID(), tailcfg.MapRequest{
+							NodeKey:   reqB.NodeKey,
+							Endpoints: []netip.AddrPort{endpoint},
+							Hostinfo: &tailcfg.Hostinfo{
+								Hostname: "current-host",
+								NetInfo:  &tailcfg.NetInfo{PreferredDERP: 7},
+							},
+						})
+						mapDone <- err
+					}()
+
+					require.Eventually(t, func() bool {
+						n, ok := s.GetNodeByID(b.node.ID())
+						return ok && n.Hostname() == "current-host"
+					}, 5*time.Second, time.Millisecond)
+
+					want.Hostname = "current-host"
+				}
+			}}
+			s.db.DB.ConnPool = barrier
+			s.db.DB.Statement.ConnPool = barrier
+
+			restore := func() {
+				s.db.DB.ConnPool = pool
+				s.db.DB.Statement.ConnPool = pool
+			}
+			defer restore()
+
+			_, err = b.reregister(b.node, pakB, reqB)
+			require.ErrorIs(t, err, types.PAKError("authkey already used"))
+
+			for _, done := range pending {
+				require.NoError(t, <-done)
+			}
+
+			restore()
+
+			ns, row := b.fields(t)
+			assert.Equal(t, want, ns, "rollback must preserve newer writes")
+			assert.Equal(t, want, row, "rollback must repair the persisted grant")
+
+			if update == "map" {
+				current, ok := s.GetNodeByID(b.node.ID())
+				require.True(t, ok)
+				require.True(t, current.Hostinfo().NetInfo().Valid())
+				assert.Equal(t, tailcfg.DERPRegionID(7), current.Hostinfo().NetInfo().PreferredDERP())
+				assert.Equal(t, []netip.AddrPort{endpoint}, current.Endpoints().AsSlice())
+
+				stored, err := s.db.GetNodeByID(b.node.ID())
+				require.NoError(t, err)
+				assert.Equal(t, "current-host", stored.Hostinfo.Hostname)
+				assert.Equal(t, tailcfg.DERPRegionID(7), stored.Hostinfo.NetInfo.PreferredDERP)
+				assert.Equal(t, []netip.AddrPort{endpoint}, stored.Endpoints.List())
+			}
+
+			require.NoError(t, s.Close())
+			reopened := persistTestReopen(t, dbPath)
+			reloaded, ok := reopened.GetNodeByID(b.node.ID())
+			require.True(t, ok)
+			assert.Equal(t, want, pakFieldsOf(reloaded.AsStruct()), "restart must keep the rollback")
+		})
+	}
+}
+
 // TestUserKeyOnTaggedMachine: a user key on a tagged machine re-authenticates
 // the tagged node only when it speaks for it (the user an admin tagged it
 // from); another user's key wakes it, gets a device of its own beside it with

hscontrol/state/state.go

@@ -17,6 +17,7 @@ import (
 	"fmt"
 	"maps"
 	"net/netip"
+	"reflect"
 	"slices"
 	"strconv"
 	"strings"
@@ -673,7 +674,6 @@ func (s *State) persistNodeWith(
 
 		return nil
 	})
-
 	if err != nil {
 		return types.NodeView{}, fmt.Errorf("saving node: %w", err)
 	}
@@ -3153,9 +3153,11 @@ func (s *State) reregisterNodeWithPAK(
 	}
 
 	// prior is the node the writer replaced: the NetInfo source and the
-	// rollback target if the database write below fails.
+	// rollback target if the database write below fails. written is what this
+	// write left; the view after the batch can hold later writes.
 	var (
 		prior   *types.Node
+		written *types.Node
 		authErr error
 		reauth  bool
 	)
@@ -3166,6 +3168,7 @@ func (s *State) reregisterNodeWithPAK(
 		now := time.Now()
 
 		prior = node.Clone()
+		defer func() { written = node.Clone() }()
 
 		reauth, authErr = pakReregister(node.View(), pak, regReq.NodeKey, now)
 		if authErr != nil {
@@ -3308,21 +3311,7 @@ func (s *State) reregisterNodeWithPAK(
 		return nil
 	})
 	if err != nil {
-		// Restore the registration fields from the node the writer replaced.
-		// Sessions, endpoints and health may have moved since that snapshot;
-		// preserve them. LastSeen stays because the node did contact us.
-		s.nodeStore.UpdateNode(prior.ID, func(n *types.Node) {
-			n.NodeKey = prior.NodeKey
-			n.Hostname = prior.Hostname
-			n.Hostinfo = prior.Hostinfo
-			n.RegisterMethod = prior.RegisterMethod
-			n.Tags = prior.Tags
-			n.UserID = prior.UserID
-			n.User = prior.User
-			n.Expiry = prior.Expiry
-			n.AuthKey = prior.AuthKey
-			n.AuthKeyID = prior.AuthKeyID
-		})
+		s.undoReregistration(prior, written)
 
 		return types.NodeView{}, false, fmt.Errorf("writing node to database: %w", err)
 	}
@@ -3339,6 +3328,56 @@ func (s *State) reregisterNodeWithPAK(
 	return updatedNodeView, onlyKeysChanged(before, updatedNodeView, time.Now()), nil
 }
 
+// undoReregistration puts back the fields of a failed re-registration that
+// still hold what it wrote, so later admin and map writes survive, and saves
+// the result. LastSeen stays because the node did contact us; sessions,
+// endpoints and health are not rolled back.
+func (s *State) undoReregistration(prior, written *types.Node) {
+	restored, exists := s.nodeStore.UpdateNode(prior.ID, func(n *types.Node) {
+		if n.NodeKey == written.NodeKey {
+			n.NodeKey = prior.NodeKey
+		}
+
+		if n.Hostname == written.Hostname {
+			n.Hostname = prior.Hostname
+		}
+
+		if reflect.DeepEqual(n.Hostinfo, written.Hostinfo) {
+			n.Hostinfo = prior.Hostinfo
+		}
+
+		if n.RegisterMethod == written.RegisterMethod {
+			n.RegisterMethod = prior.RegisterMethod
+		}
+
+		if slices.Equal(n.Tags, written.Tags) &&
+			reflect.DeepEqual(n.UserID, written.UserID) &&
+			reflect.DeepEqual(n.User, written.User) {
+			n.Tags = prior.Tags
+			n.UserID = prior.UserID
+			n.User = prior.User
+		}
+
+		if reflect.DeepEqual(n.Expiry, written.Expiry) {
+			n.Expiry = prior.Expiry
+		}
+
+		if reflect.DeepEqual(n.AuthKeyID, written.AuthKeyID) &&
+			reflect.DeepEqual(n.AuthKey, written.AuthKey) {
+			n.AuthKey = prior.AuthKey
+			n.AuthKeyID = prior.AuthKeyID
+		}
+	})
+	if !exists {
+		return
+	}
+
+	_, err := s.persistNode(restored)
+	if err != nil {
+		log.Error().Err(err).Uint64(zf.NodeID, prior.ID.Uint64()).Msg("Persisting re-registration rollback failed")
+	}
+}
+
 // reauthChange returns the [change.Change] to broadcast after an authentication
 // that updated or created a node. keysOnly marks a same-user relogin that
 // changed nothing peers read besides its keys; see [onlyKeysChanged].