implement followup part 2

Signed-off-by: Kristoffer Dalby <[email protected]>

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

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
 )