From d85c3ae555e5025da33d10566bc04bc6987ffac4 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Mon, 13 Jan 2025 11:21:38 +0100
Subject: [PATCH] implement followup part 2

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
---
 hscontrol/auth.go     | 70 +++++++++++++++++++++++--------------------
 hscontrol/handlers.go |  3 --
 2 files changed, 38 insertions(+), 35 deletions(-)

diff --git a/hscontrol/auth.go b/hscontrol/auth.go
index 487bb9f0ee..b0d632b9e1 100644
--- a/hscontrol/auth.go
+++ b/hscontrol/auth.go
@@ -69,6 +69,40 @@ func logAuthFunc(
 		}
 }
 
+func (h *Headscale) waitForFollowup(
+	req *http.Request,
+	regReq tailcfg.RegisterRequest,
+	logTrace func(string),
+) {
+	logTrace("register request is a followup")
+	fu, err := url.Parse(regReq.Followup)
+	if err != nil {
+		logTrace("failed to parse followup URL")
+		return
+	}
+
+	followupReg, err := types.RegistrationIDFromString(strings.ReplaceAll(fu.Path, "/register/", ""))
+	if err != nil {
+		logTrace("followup URL does not contains a valid registration ID")
+		return
+	}
+
+	logTrace(fmt.Sprintf("followup URL contains a valid registration ID, looking up in cache: %s", followupReg))
+
+	if reg, ok := h.registrationCache.Get(followupReg); ok {
+		logTrace("Node is waiting for interactive login")
+
+		select {
+		case <-req.Context().Done():
+			logTrace("node went away before it was registered")
+			return
+		case <-reg.Registered:
+			logTrace("node has successfully registered")
+			return
+		}
+	}
+}
+
 // handleRegister is the logic for registering a client.
 func (h *Headscale) handleRegister(
 	writer http.ResponseWriter,
@@ -102,33 +136,8 @@ func (h *Headscale) handleRegister(
 
 		// Check if the node is waiting for interactive login.
 		if regReq.Followup != "" {
-			logTrace("register request is a followup")
-			fu, err := url.Parse(regReq.Followup)
-			if err != nil {
-				logTrace("failed to parse followup URL")
-				return
-			}
-
-			followupReg, err := types.RegistrationIDFromString(strings.ReplaceAll(fu.Path, "/register/", ""))
-			if err != nil {
-				logTrace("followup URL does not contains a valid registration ID")
-				return
-			}
-
-			logTrace(fmt.Sprintf("followup URL contains a valid registration ID, looking up in cache: %s", followupReg))
-
-			if reg, ok := h.registrationCache.Get(followupReg); ok {
-				logTrace("Node is waiting for interactive login")
-
-				select {
-				case <-req.Context().Done():
-					logTrace("node went away before it was successfully registered")
-					return
-				case <-reg.Registered:
-					logTrace("node has successfully registered")
-					return
-				}
-			}
+			h.waitForFollowup(req, regReq, logTrace)
+			return
 		}
 
 		logInfo("Node not found in database, creating new")
@@ -233,11 +242,8 @@ func (h *Headscale) handleRegister(
 		}
 
 		if regReq.Followup != "" {
-			select {
-			case <-req.Context().Done():
-				return
-			case <-time.After(registrationHoldoff):
-			}
+			h.waitForFollowup(req, regReq, logTrace)
+			return
 		}
 
 		// The node has expired or it is logged out
diff --git a/hscontrol/handlers.go b/hscontrol/handlers.go
index e4a3857c5c..3b7829a96f 100644
--- a/hscontrol/handlers.go
+++ b/hscontrol/handlers.go
@@ -8,7 +8,6 @@ import (
 	"net/http"
 	"strconv"
 	"strings"
-	"time"
 
 	"github.com/chasefleming/elem-go"
 	"github.com/chasefleming/elem-go/attrs"
@@ -33,8 +32,6 @@ const (
 	// See also https://github.com/tailscale/tailscale/blob/main/tailcfg/tailcfg.go
 	NoiseCapabilityVersion = 39
 
-	// TODO(juan): remove this once https://github.com/juanfont/headscale/issues/727 is fixed.
-	registrationHoldoff        = time.Second * 5
 	reservedResponseHeaderSize = 4
 )