From 2504f0de4c22f934e2ec17e6ce5ba74475adf4a1 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Tue, 25 Feb 2025 10:33:20 +0100
Subject: [PATCH] return regresp or err after waiting for registration

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
---
 hscontrol/auth.go         | 32 +++++++++++++++++++-------------
 hscontrol/db/node.go      |  5 +++++
 hscontrol/grpcv1.go       |  2 +-
 hscontrol/types/common.go |  2 +-
 4 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/hscontrol/auth.go b/hscontrol/auth.go
index 4cc7058b..0a8602cd 100644
--- a/hscontrol/auth.go
+++ b/hscontrol/auth.go
@@ -43,10 +43,7 @@ func (h *Headscale) handleRegister(
 	}
 
 	if regReq.Followup != "" {
-		// TODO(kradalby): Does this need to return an error of some sort?
-		// Maybe if the registration fails down the line it can be sent
-		// on the channel and returned here?
-		h.waitForFollowup(ctx, regReq)
+		return h.waitForFollowup(ctx, regReq)
 	}
 
 	if regReq.Auth != nil && regReq.Auth.AuthKey != "" {
@@ -111,42 +108,51 @@ func (h *Headscale) handleExistingNode(
 		h.nodeNotifier.NotifyWithIgnore(ctx, types.UpdateExpire(node.ID, requestExpiry), node.ID)
 	}
 
+	return nodeToRegisterResponse(node), nil
+}
+
+func nodeToRegisterResponse(node *types.Node) *tailcfg.RegisterResponse {
 	return &tailcfg.RegisterResponse{
 		// TODO(kradalby): Only send for user-owned nodes
 		// and not tagged nodes when tags is working.
 		User:           *node.User.TailscaleUser(),
 		Login:          *node.User.TailscaleLogin(),
-		NodeKeyExpired: expired,
+		NodeKeyExpired: node.IsExpired(),
 
 		// Headscale does not implement the concept of machine authorization
 		// so we always return true here.
 		// Revisit this if #2176 gets implemented.
 		MachineAuthorized: true,
-	}, nil
+	}
 }
 
 func (h *Headscale) waitForFollowup(
 	ctx context.Context,
 	regReq tailcfg.RegisterRequest,
-) {
+) (*tailcfg.RegisterResponse, error) {
 	fu, err := url.Parse(regReq.Followup)
 	if err != nil {
-		return
+		return nil, NewHTTPError(http.StatusUnauthorized, "invalid followup URL", err)
 	}
 
 	followupReg, err := types.RegistrationIDFromString(strings.ReplaceAll(fu.Path, "/register/", ""))
 	if err != nil {
-		return
+		return nil, NewHTTPError(http.StatusUnauthorized, "invalid registration ID", err)
 	}
 
 	if reg, ok := h.registrationCache.Get(followupReg); ok {
 		select {
 		case <-ctx.Done():
-			return
-		case <-reg.Registered:
-			return
+			return nil, NewHTTPError(http.StatusUnauthorized, "registration timed out", err)
+		case node := <-reg.Registered:
+			if node == nil {
+				return nil, NewHTTPError(http.StatusUnauthorized, "node not found", nil)
+			}
+			return nodeToRegisterResponse(node), nil
 		}
 	}
+
+	return nil, NewHTTPError(http.StatusNotFound, "followup registration not found", nil)
 }
 
 // canUsePreAuthKey checks if a pre auth key can be used.
@@ -271,7 +277,7 @@ func (h *Headscale) handleRegisterInteractive(
 			Hostinfo:   regReq.Hostinfo,
 			LastSeen:   ptr.To(time.Now()),
 		},
-		Registered: make(chan struct{}),
+		Registered: make(chan *types.Node),
 	}
 
 	if !regReq.Expiry.IsZero() {
diff --git a/hscontrol/db/node.go b/hscontrol/db/node.go
index c9244095..74cd7a9f 100644
--- a/hscontrol/db/node.go
+++ b/hscontrol/db/node.go
@@ -372,7 +372,12 @@ func (hsdb *HSDatabase) HandleNodeFromAuthPath(
 				}
 
 				// Signal to waiting clients that the machine has been registered.
+				select {
+				case reg.Registered <- node:
+				default:
+				}
 				close(reg.Registered)
+
 				newNode = true
 				return node, err
 			} else {
diff --git a/hscontrol/grpcv1.go b/hscontrol/grpcv1.go
index 59fe4ebd..7368083c 100644
--- a/hscontrol/grpcv1.go
+++ b/hscontrol/grpcv1.go
@@ -838,7 +838,7 @@ func (api headscaleV1APIServer) DebugCreateNode(
 
 			Hostinfo: &hostinfo,
 		},
-		Registered: make(chan struct{}),
+		Registered: make(chan *types.Node),
 	}
 
 	log.Debug().
diff --git a/hscontrol/types/common.go b/hscontrol/types/common.go
index c8d696af..c4cc8a2e 100644
--- a/hscontrol/types/common.go
+++ b/hscontrol/types/common.go
@@ -194,5 +194,5 @@ func (r RegistrationID) String() string {
 
 type RegisterNode struct {
 	Node       Node
-	Registered chan struct{}
+	Registered chan *Node
 }