[management] batch disconnect and lock free reconnect (#8373)

management/server/peer_test.go

@@ -17,6 +17,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -3323,3 +3324,136 @@ func TestPeerWrites_BumpNetworkSerial(t *testing.T) {
 		assert.Greater(t, networkSerial(t, manager, account.Id), afterDisconnect, "embedded proxy connect must bump")
 	})
 }
+
+type txCountingStore struct {
+	store.Store
+	txs atomic.Int32
+}
+
+func (s *txCountingStore) ExecuteInTransaction(ctx context.Context, f func(store.Store) error) error {
+	s.txs.Add(1)
+	return s.Store.ExecuteInTransaction(ctx, f)
+}
+
+func TestPeerSyncAndLogin_LockOnlyOnChange(t *testing.T) {
+	manager, _, account, peer1, _, _ := setupNetworkMapTest(t)
+	ctx := context.Background()
+	counting := &txCountingStore{Store: manager.Store}
+	manager.Store = counting
+
+	current, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+	require.NoError(t, err)
+
+	t.Run("SyncPeer without meta change opens no transaction", func(t *testing.T) {
+		counting.txs.Store(0)
+		_, _, _, _, err := manager.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peer1.Key, Meta: current.Meta}, account.Id)
+		require.NoError(t, err)
+		assert.Zero(t, counting.txs.Load())
+	})
+
+	t.Run("SyncPeer with meta change saves in one transaction", func(t *testing.T) {
+		counting.txs.Store(0)
+		meta := current.Meta
+		meta.Hostname = "lock-on-change-host"
+		_, _, _, _, err := manager.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peer1.Key, Meta: meta}, account.Id)
+		require.NoError(t, err)
+		assert.Equal(t, int32(1), counting.txs.Load())
+
+		saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+		require.NoError(t, err)
+		assert.Equal(t, "lock-on-change-host", saved.Meta.Hostname)
+	})
+
+	t.Run("LoginPeer without change opens no transaction", func(t *testing.T) {
+		saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+		require.NoError(t, err)
+
+		counting.txs.Store(0)
+		_, _, _, _, err = manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, SSHKey: saved.SSHKey, Meta: saved.Meta})
+		require.NoError(t, err)
+		assert.Zero(t, counting.txs.Load())
+	})
+
+	t.Run("LoginPeer with new ssh key saves in one transaction", func(t *testing.T) {
+		counting.txs.Store(0)
+		_, _, _, _, err := manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, SSHKey: "ssh-ed25519 AAAAlockonchange", Meta: current.Meta})
+		require.NoError(t, err)
+		assert.Equal(t, int32(1), counting.txs.Load())
+
+		saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+		require.NoError(t, err)
+		assert.Equal(t, "ssh-ed25519 AAAAlockonchange", saved.SSHKey)
+	})
+
+	t.Run("LoginPeer rejects a user mismatch without a transaction", func(t *testing.T) {
+		counting.txs.Store(0)
+		_, _, _, _, err := manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, UserID: "someone-else", SSHKey: "ssh-ed25519 AAAAother", Meta: current.Meta})
+		assert.Equal(t, status.NewPeerLoginMismatchError(), err)
+		assert.Zero(t, counting.txs.Load())
+
+		saved, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+		require.NoError(t, err)
+		assert.Equal(t, "ssh-ed25519 AAAAlockonchange", saved.SSHKey, "a rejected login must not store its ssh key")
+	})
+
+	t.Run("LoginPeer rejects disallowed extra DNS labels without a transaction", func(t *testing.T) {
+		counting.txs.Store(0)
+		_, _, _, _, err := manager.LoginPeer(ctx, types.PeerLogin{WireGuardPubKey: peer1.Key, SSHKey: "ssh-ed25519 AAAAlabels", ExtraDNSLabels: []string{"extra"}, Meta: current.Meta})
+		errStatus, ok := status.FromError(err)
+		require.True(t, ok, "expected a status error, got %v", err)
+		assert.Equal(t, status.PreconditionFailed, errStatus.Type())
+		assert.Zero(t, counting.txs.Load())
+	})
+}
+
+func TestCheckPeerMaySync(t *testing.T) {
+	manager, _, account, peer1, peer2, _ := setupNetworkMapTest(t)
+	ctx := context.Background()
+
+	require.NoError(t, manager.CheckPeerMaySync(ctx, account.Id, peer1.Key, nil))
+
+	expired, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer1.ID)
+	require.NoError(t, err)
+	expired.Status.LoginExpired = true
+	require.NoError(t, manager.Store.SavePeer(ctx, account.Id, expired))
+	assert.Equal(t, status.NewPeerLoginExpiredError(), manager.CheckPeerMaySync(ctx, account.Id, peer1.Key, nil))
+
+	owner := types.NewRegularUser("blocked-owner", "", "")
+	owner.AccountID = account.Id
+	owner.Blocked = true
+	require.NoError(t, manager.Store.SaveUser(ctx, owner))
+	blocked, err := manager.Store.GetPeerByID(ctx, store.LockingStrengthNone, account.Id, peer2.ID)
+	require.NoError(t, err)
+	blocked.UserID = owner.Id
+	require.NoError(t, manager.Store.SavePeer(ctx, account.Id, blocked))
+	err = manager.CheckPeerMaySync(ctx, account.Id, peer2.Key, nil)
+	errStatus, ok := status.FromError(err)
+	require.True(t, ok, "expected a status error, got %v", err)
+	assert.Equal(t, status.PermissionDenied, errStatus.Type())
+
+	assert.Equal(t, status.NewPeerNotRegisteredError(), manager.CheckPeerMaySync(ctx, account.Id, "unknown-key", nil))
+}
+
+type failingPeerLookupStore struct {
+	store.Store
+}
+
+func (s *failingPeerLookupStore) GetPeerByPeerPubKey(context.Context, store.LockingStrength, string) (*nbpeer.Peer, error) {
+	return nil, status.Errorf(status.Internal, "issue getting peer from store: connection reset")
+}
+
+func TestPeerLookupErrorKeepsTransientStoreErrors(t *testing.T) {
+	manager, _, account, peer1, _, _ := setupNetworkMapTest(t)
+	ctx := context.Background()
+	manager.Store = &failingPeerLookupStore{Store: manager.Store}
+
+	err := manager.CheckPeerMaySync(ctx, account.Id, peer1.Key, nil)
+	errStatus, ok := status.FromError(err)
+	require.True(t, ok, "expected a status error, got %v", err)
+	assert.Equal(t, status.Internal, errStatus.Type(), "a store failure must not look like an unregistered peer")
+
+	_, _, _, _, err = manager.SyncPeer(ctx, types.PeerSync{WireGuardPubKey: peer1.Key}, account.Id)
+	errStatus, ok = status.FromError(err)
+	require.True(t, ok, "expected a status error, got %v", err)
+	assert.Equal(t, status.Internal, errStatus.Type(), "a store failure must not look like an unregistered peer")
+}