wgengine/userspace: add extra check for tsmp learned keys in engine (#19223)

If an entry in the tsmpLearnedDisco does not match the disco key of the
key currently being processed, overwrite the key, and leave the entry in
the map for later processing.

In reality, this should not happen, but is put in as a safety measure
with logging of the situation so we can replicate the behaviour and
correct it should it happen.

Updates #12639

Signed-off-by: Claus Lensbøl <claus@tailscale.com>

wgengine/userspace_test.go

@@ -237,6 +237,93 @@ func TestUserspaceEngineTSMPLearned(t *testing.T) {
 	}
 }
 
+func TestUserspaceEngineTSMPLearnedMismatch(t *testing.T) {
+	bus := eventbustest.NewBus(t)
+
+	ht := health.NewTracker(bus)
+	reg := new(usermetric.Registry)
+	e, err := NewFakeUserspaceEngine(t.Logf, 0, ht, reg, bus)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(e.Close)
+	ue := e.(*userspaceEngine)
+
+	discoChangedChan := make(chan map[key.NodePublic]bool, 1)
+	ue.testDiscoChangedHook = func(m map[key.NodePublic]bool) {
+		discoChangedChan <- m
+	}
+
+	routerCfg := &router.Config{}
+	var metricValue int64 = 0
+
+	keyChanges := []struct {
+		tsmp     bool
+		inMap    bool
+		wrongKey bool
+	}{
+		{tsmp: false, inMap: false, wrongKey: false},
+		{tsmp: true, inMap: false, wrongKey: true},
+		{tsmp: false, inMap: false, wrongKey: false},
+	}
+
+	nkHex := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
+	for _, change := range keyChanges {
+		oldDisco := key.NewDisco()
+		nm := &netmap.NetworkMap{
+			Peers: nodeViews([]*tailcfg.Node{
+				{
+					ID:       1,
+					Key:      nkFromHex(nkHex),
+					DiscoKey: oldDisco.Public(),
+				},
+			}),
+		}
+		nk, err := key.ParseNodePublicUntyped(mem.S(nkHex))
+		if err != nil {
+			t.Fatal(err)
+		}
+		e.SetNetworkMap(nm)
+
+		newDisco := key.NewDisco()
+		cfg := &wgcfg.Config{
+			Peers: []wgcfg.Peer{
+				{
+					PublicKey: nk,
+					DiscoKey:  newDisco.Public(),
+				},
+			},
+		}
+
+		tsmpKey := newDisco.Public()
+		if change.tsmp {
+			if change.wrongKey {
+				tsmpKey = key.NewDisco().Public()
+			}
+			ue.PatchDiscoKey(nk, tsmpKey)
+		}
+		err = e.Reconfig(cfg, routerCfg, &dns.Config{})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		changeMap := <-discoChangedChan
+
+		if _, ok := changeMap[nk]; ok != change.inMap {
+			t.Fatalf("expect key %v in map %v to be %t, got %t", nk, changeMap,
+				change.inMap, ok)
+		}
+
+		metric := metricTSMPLearnedKeyMismatch.Value()
+		delta := metric - metricValue
+		metricValue = metric
+
+		if change.wrongKey && delta != 1 {
+			t.Fatalf("expected a delta of 1, got %d", delta)
+		}
+	}
+}
+
 func TestUserspaceEnginePortReconfig(t *testing.T) {
 	flakytest.Mark(t, "https://github.com/tailscale/tailscale/issues/2855")
 	const defaultPort = 49983