From d972099c78dab2327fed94536750e6233084c846 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Wed, 10 Nov 2021 12:09:54 -0800
Subject: [PATCH] wgengine/magicsock: add a stress test

Updates tailscale/corp#3016

Change-Id: I23708e68ed44d81986d9e2be82029d4555547592
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 wgengine/magicsock/magicsock_test.go | 57 ++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/wgengine/magicsock/magicsock_test.go b/wgengine/magicsock/magicsock_test.go
index 0872c9c7c..7b1ac3ff3 100644
--- a/wgengine/magicsock/magicsock_test.go
+++ b/wgengine/magicsock/magicsock_test.go
@@ -12,6 +12,7 @@ import (
 	"errors"
 	"fmt"
 	"io/ioutil"
+	"math/rand"
 	"net"
 	"net/http"
 	"net/http/httptest"
@@ -1639,3 +1640,59 @@ func epStrings(eps []tailcfg.Endpoint) (ret []string) {
 	}
 	return
 }
+
+func TestStressSetNetworkMap(t *testing.T) {
+	conn := newTestConn(t)
+	t.Cleanup(func() { conn.Close() })
+	var buf tstest.MemLogger
+	conn.logf = buf.Logf
+
+	conn.SetPrivateKey(wgkey.Private{0: 1})
+
+	const num = 5
+	present := make([]bool, num)
+	allPeers := make([]*tailcfg.Node, 0, num)
+	for i := 0; i < num; i++ {
+		present[i] = true
+		allPeers = append(allPeers, &tailcfg.Node{
+			DiscoKey:  randDiscoKey(),
+			Key:       randNodeKey(),
+			Endpoints: []string{"192.168.1.2:345"},
+		})
+	}
+
+	var peers []*tailcfg.Node
+	for i := 0; i < 1000; i++ {
+		which := rand.Intn(num)
+		action := rand.Intn(3)
+		switch action {
+		case 0:
+			present[which] = !present[which]
+		case 1:
+			allPeers[which].DiscoKey = randDiscoKey()
+		case 2:
+			allPeers[which].Key = randNodeKey()
+		default:
+			panic("unreachable")
+		}
+		peers = peers[:0]
+		for peerIdx, p := range allPeers {
+			if present[peerIdx] {
+				peers = append(peers, p)
+			}
+		}
+		conn.SetNetworkMap(&netmap.NetworkMap{
+			Peers: peers,
+		})
+	}
+}
+
+func randDiscoKey() (k tailcfg.DiscoKey) {
+	crand.Read(k[:])
+	return
+}
+
+func randNodeKey() (k tailcfg.NodeKey) {
+	crand.Read(k[:])
+	return
+}