From 546f1f0a41c0dca1f688f692d906f392d8f5bfd3 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Tue, 29 Aug 2023 08:28:51 +0200
Subject: [PATCH] temp

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
---
 Dockerfile.tailscale-HEAD | 3 ++-
 hscontrol/app.go          | 5 ++++-
 hscontrol/auth.go         | 6 ++++++
 hscontrol/auth_noise.go   | 3 +++
 hscontrol/poll.go         | 6 +++---
 hscontrol/poll_noise.go   | 1 +
 integration/tsic/tsic.go  | 2 +-
 7 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/Dockerfile.tailscale-HEAD b/Dockerfile.tailscale-HEAD
index 2a3aac769c..f17a50b27b 100644
--- a/Dockerfile.tailscale-HEAD
+++ b/Dockerfile.tailscale-HEAD
@@ -6,7 +6,8 @@ RUN apt-get update \
 
 RUN useradd --shell=/bin/bash --create-home ssh-it-user
 
-RUN git clone https://github.com/tailscale/tailscale.git
+# RUN git clone https://github.com/tailscale/tailscale.git
+COPY tailscale /go/tailscale
 
 WORKDIR /go/tailscale
 
diff --git a/hscontrol/app.go b/hscontrol/app.go
index 57fb5848ff..9447fad284 100644
--- a/hscontrol/app.go
+++ b/hscontrol/app.go
@@ -49,6 +49,8 @@ import (
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/dnstype"
 	"tailscale.com/types/key"
+
+	_ "net/http/pprof"
 )
 
 var (
@@ -93,7 +95,7 @@ type Headscale struct {
 	shutdownChan       chan struct{}
 	pollNetMapStreamWG sync.WaitGroup
 
-	pollStreamOpenMu sync.Mutex
+	// pollStreamOpenMu sync.Mutex
 }
 
 func NewHeadscale(cfg *types.Config) (*Headscale, error) {
@@ -439,6 +441,7 @@ func (h *Headscale) ensureUnixSocketIsAbsent() error {
 
 func (h *Headscale) createRouter(grpcMux *runtime.ServeMux) *mux.Router {
 	router := mux.NewRouter()
+	router.PathPrefix("/debug/pprof/").Handler(http.DefaultServeMux)
 
 	router.HandleFunc(ts2021UpgradePath, h.NoiseUpgradeHandler).Methods(http.MethodPost)
 
diff --git a/hscontrol/auth.go b/hscontrol/auth.go
index e9121203c5..fe71df12b0 100644
--- a/hscontrol/auth.go
+++ b/hscontrol/auth.go
@@ -27,10 +27,14 @@ func (h *Headscale) handleRegister(
 	isNoise bool,
 ) {
 	now := time.Now().UTC()
+	log.Trace().Msg("STEP 2")
 	machine, err := h.db.GetMachineByAnyKey(machineKey, registerRequest.NodeKey, registerRequest.OldNodeKey)
+	log.Trace().Msg("STEP 3")
 	if errors.Is(err, gorm.ErrRecordNotFound) {
+		log.Trace().Msg("STEP 4")
 		// If the machine has AuthKey set, handle registration via PreAuthKeys
 		if registerRequest.Auth.AuthKey != "" {
+			log.Trace().Msg("STEP 5")
 			h.handleAuthKey(writer, registerRequest, machineKey, isNoise)
 
 			return
@@ -236,7 +240,9 @@ func (h *Headscale) handleAuthKey(
 		Msgf("Processing auth key for %s", registerRequest.Hostinfo.Hostname)
 	resp := tailcfg.RegisterResponse{}
 
+	log.Trace().Msg("STEP 5")
 	pak, err := h.db.ValidatePreAuthKey(registerRequest.Auth.AuthKey)
+	log.Trace().Msg("STEP 6")
 	if err != nil {
 		log.Error().
 			Caller().
diff --git a/hscontrol/auth_noise.go b/hscontrol/auth_noise.go
index 7f6a7fd560..25d9db3654 100644
--- a/hscontrol/auth_noise.go
+++ b/hscontrol/auth_noise.go
@@ -23,6 +23,7 @@ func (ns *noiseServer) NoiseRegistrationHandler(
 
 	log.Trace().
 		Any("headers", req.Header).
+		Caller().
 		Msg("Headers")
 
 	body, _ := io.ReadAll(req.Body)
@@ -40,5 +41,7 @@ func (ns *noiseServer) NoiseRegistrationHandler(
 
 	ns.nodeKey = registerRequest.NodeKey
 
+	log.Trace().Msg("STEP 1")
+
 	ns.headscale.handleRegister(writer, req, registerRequest, ns.conn.Peer(), true)
 }
diff --git a/hscontrol/poll.go b/hscontrol/poll.go
index dc9763b10f..462207fd57 100644
--- a/hscontrol/poll.go
+++ b/hscontrol/poll.go
@@ -69,7 +69,7 @@ func (h *Headscale) handlePoll(
 	// following updates missing
 	var updateChan chan types.StateUpdate
 	if mapRequest.Stream {
-		h.pollStreamOpenMu.Lock()
+		// h.pollStreamOpenMu.Lock()
 		h.pollNetMapStreamWG.Add(1)
 		defer h.pollNetMapStreamWG.Done()
 
@@ -233,6 +233,8 @@ func (h *Headscale) handlePoll(
 		return
 	}
 
+	// h.pollStreamOpenMu.Unlock()
+
 	// Send the client an update to make sure we send an initial mapresponse
 	_, err = writer.Write(mapResp)
 	if err != nil {
@@ -254,8 +256,6 @@ func (h *Headscale) handlePoll(
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	h.pollStreamOpenMu.Unlock()
-
 	for {
 		logInfo("Waiting for update on stream channel")
 		select {
diff --git a/hscontrol/poll_noise.go b/hscontrol/poll_noise.go
index a21ef5ff88..3d672f0b6b 100644
--- a/hscontrol/poll_noise.go
+++ b/hscontrol/poll_noise.go
@@ -31,6 +31,7 @@ func (ns *noiseServer) NoisePollNetMapHandler(
 
 	log.Trace().
 		Any("headers", req.Header).
+		Caller().
 		Msg("Headers")
 
 	body, _ := io.ReadAll(req.Body)
diff --git a/integration/tsic/tsic.go b/integration/tsic/tsic.go
index 08d7447514..108ce5f06d 100644
--- a/integration/tsic/tsic.go
+++ b/integration/tsic/tsic.go
@@ -340,7 +340,7 @@ func (t *TailscaleInContainer) Login(
 		)
 	}
 
-	if _, _, err := t.Execute(command); err != nil {
+	if _, _, err := t.Execute(command, dockertestutil.ExecuteCommandTimeout(60*time.Second)); err != nil {
 		return fmt.Errorf(
 			"failed to join tailscale client (%s): %w",
 			strings.Join(command, " "),