From 78268d78a01a34c0884a52af954cf3bfe855a8be Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Mon, 17 Jul 2023 11:13:48 +0200
Subject: [PATCH] add debug option to save all map responses

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
---
 hscontrol/mapper/mapper.go | 42 ++++++++++++++++++++++++++++++++++++++
 hscontrol/poll.go          | 11 ++++++----
 integration/hsic/hsic.go   | 27 ++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 4 deletions(-)

diff --git a/hscontrol/mapper/mapper.go b/hscontrol/mapper/mapper.go
index 819a7fb4..5f88b839 100644
--- a/hscontrol/mapper/mapper.go
+++ b/hscontrol/mapper/mapper.go
@@ -4,10 +4,14 @@ import (
 	"encoding/binary"
 	"encoding/json"
 	"fmt"
+	"io/fs"
 	"net/url"
+	"os"
+	"path"
 	"sort"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	mapset "github.com/deckarep/golang-set/v2"
@@ -18,6 +22,7 @@ import (
 	"github.com/klauspost/compress/zstd"
 	"github.com/rs/zerolog/log"
 	"github.com/samber/lo"
+	"tailscale.com/envknob"
 	"tailscale.com/smallzstd"
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/dnstype"
@@ -29,6 +34,8 @@ const (
 	reservedResponseHeaderSize = 4
 )
 
+var debugDumpMapResponsePath = envknob.String("HEADSCALE_DEBUG_DUMP_MAPRESPONSE_PATH")
+
 type Mapper struct {
 	db *db.HSDatabase
 
@@ -413,6 +420,41 @@ func (m Mapper) marshalMapResponse(
 			Msg("Cannot marshal map response")
 	}
 
+	if debugDumpMapResponsePath != "" {
+		data := map[string]interface{}{
+			"MapRequest":  mapRequest,
+			"MapResponse": resp,
+		}
+
+		body, err := json.Marshal(data)
+		if err != nil {
+			log.Error().
+				Caller().
+				Err(err).
+				Msg("Cannot marshal map response")
+		}
+
+		perms := fs.FileMode(debugMapResponsePerm)
+		mPath := path.Join(debugDumpMapResponsePath, machine.Hostname)
+		err = os.MkdirAll(mPath, perms)
+		if err != nil {
+			panic(err)
+		}
+
+		now := time.Now().Unix()
+
+		mapResponsePath := path.Join(
+			mPath,
+			fmt.Sprintf("%d-%s-%d.json", now, m.uid, atomic.LoadUint64(&m.seq)),
+		)
+
+		log.Trace().Msgf("Writing MapResponse to %s", mapResponsePath)
+		err = os.WriteFile(mapResponsePath, body, perms)
+		if err != nil {
+			panic(err)
+		}
+	}
+
 	var respBody []byte
 	if compression == util.ZstdCompression {
 		respBody = zstdEncode(jsonBody)
diff --git a/hscontrol/poll.go b/hscontrol/poll.go
index 5197e731..bf7a0f49 100644
--- a/hscontrol/poll.go
+++ b/hscontrol/poll.go
@@ -225,8 +225,7 @@ func (h *Headscale) pollNetMapStream(
 	h.pollNetMapStreamWG.Add(1)
 	defer h.pollNetMapStreamWG.Done()
 
-	const chanSize = 8
-	updateChan := make(chan types.StateUpdate, chanSize)
+	updateChan := make(chan types.StateUpdate)
 	defer closeChanWithLog(updateChan, machine.Hostname, "updateChan")
 
 	// Register the node's update channel
@@ -271,14 +270,18 @@ func (h *Headscale) pollNetMapStream(
 			var err error
 
 			switch update.Type {
-			case types.StateFullUpdate:
-				data, err = mapp.FullMapResponse(mapRequest, machine, h.ACLPolicy)
 			case types.StatePeerChanged:
+				logInfo("Sending PeerChanged MapResponse")
 				data, err = mapp.PeerChangedResponse(mapRequest, machine, update.Changed, h.ACLPolicy)
 			case types.StatePeerRemoved:
+				logInfo("Sending PeerRemoved MapResponse")
 				data, err = mapp.PeerRemovedResponse(mapRequest, machine, update.Removed)
 			case types.StateDERPUpdated:
+				logInfo("Sending DERPUpdate MapResponse")
 				data, err = mapp.DERPMapResponse(mapRequest, machine, update.DERPMap)
+			case types.StateFullUpdate:
+				logInfo("Sending Full MapResponse")
+				data, err = mapp.FullMapResponse(mapRequest, machine, h.ACLPolicy)
 			}
 
 			if err != nil {
diff --git a/integration/hsic/hsic.go b/integration/hsic/hsic.go
index e13b7273..01f34443 100644
--- a/integration/hsic/hsic.go
+++ b/integration/hsic/hsic.go
@@ -212,6 +212,7 @@ func New(
 	env := []string{
 		"HEADSCALE_PROFILING_ENABLED=1",
 		"HEADSCALE_PROFILING_PATH=/tmp/profile",
+		"HEADSCALE_DEBUG_DUMP_MAPRESPONSE_PATH=/tmp/mapresponses",
 	}
 	for key, value := range hsic.env {
 		env = append(env, fmt.Sprintf("%s=%s", key, value))
@@ -339,6 +340,14 @@ func (t *HeadscaleInContainer) Shutdown() error {
 		)
 	}
 
+	err = t.SaveMapResponses("/tmp/control")
+	if err != nil {
+		log.Printf(
+			"Failed to save mapresponses from control: %s",
+			fmt.Errorf("failed to save mapresponses from control: %w", err),
+		)
+	}
+
 	return t.pool.Purge(t.container)
 }
 
@@ -354,6 +363,24 @@ func (t *HeadscaleInContainer) SaveProfile(savePath string) error {
 		return err
 	}
 
+	err = os.WriteFile(
+		path.Join(savePath, t.hostname+"maps.tar"),
+		tarFile,
+		os.ModePerm,
+	)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (t *HeadscaleInContainer) SaveMapResponses(savePath string) error {
+	tarFile, err := t.FetchPath("/tmp/mapresponses")
+	if err != nil {
+		return err
+	}
+
 	err = os.WriteFile(
 		path.Join(savePath, t.hostname+".pprof.tar"),
 		tarFile,