refactor: some renaming and restructuring

hscontrol/app.go

@@ -219,9 +219,16 @@ func NewHeadscale(cfg *types.Config) (*Headscale, error) {
 			)
 		}
 
+		if cfg.DERP.ServerVerifyClients {
+			t := http.DefaultTransport.(*http.Transport)
+			t.RegisterProtocol(
+				derpServer.DerpVerifyScheme,
+				derpServer.NewDERPVerifyTransport(app.handleVerifyRequest),
+			)
+		}
+
 		embeddedDERPServer, err := derpServer.NewDERPServer(
 			cfg.ServerURL,
-			app.VerifyHandler,
 			key.NodePrivate(*derpServerKey),
 			&cfg.DERP,
 		)

hscontrol/derp/server/derp_server.go

@@ -2,12 +2,13 @@ package server
 
 import (
 	"bufio"
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
-	"net/http/httptest"
 	"net/netip"
 	"net/url"
 	"strconv"
@@ -30,6 +31,7 @@ import (
 // headers and it will begin writing & reading the DERP protocol immediately
 // following its HTTP request.
 const fastStartHeader = "Derp-Fast-Start"
+const DerpVerifyScheme = "derp-verify"
 
 type DERPServer struct {
 	serverURL     string
@@ -40,7 +42,6 @@ type DERPServer struct {
 
 func NewDERPServer(
 	serverURL string,
-	verifyHandler http.HandlerFunc,
 	derpKey key.NodePrivate,
 	cfg *types.DERPConfig,
 ) (*DERPServer, error) {
@@ -48,11 +49,7 @@ func NewDERPServer(
 	server := derp.NewServer(derpKey, util.TSLogfWrapper()) // nolint // zerolinter complains
 
 	if cfg.ServerVerifyClients {
-		t := http.DefaultTransport.(*http.Transport)
+		server.SetVerifyClientURL(DerpVerifyScheme + "://verify")
-		t.RegisterProtocol("headscale", &HeadscaleTransport{
-			verifyHandler: verifyHandler,
-		})
-		server.SetVerifyClientURL("headscale://verify")
 		server.SetVerifyClientURLFailOpen(false)
 	}
 
@@ -372,13 +369,31 @@ func serverSTUNListener(ctx context.Context, packetConn *net.UDPConn) {
 	}
 }
 
-type HeadscaleTransport struct {
+func NewDERPVerifyTransport(handleVerifyRequest func(*http.Request, io.Writer) error) *DERPVerifyTransport {
-	verifyHandler http.HandlerFunc
+	return &DERPVerifyTransport{
+		handleVerifyRequest: handleVerifyRequest,
+	}
+}
+
+type DERPVerifyTransport struct {
+	handleVerifyRequest func(*http.Request, io.Writer) error
+}
+
+func (t *DERPVerifyTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	buf := new(bytes.Buffer)
+	if err := t.handleVerifyRequest(req, buf); err != nil {
+		log.Error().
+			Caller().
+			Err(err).
+			Msg("Failed to handle verify request")
+
+		return nil, err
+	}
+
+	resp := &http.Response{
+		StatusCode: http.StatusOK,
+		Body:       io.NopCloser(buf),
 	}
 
-func (t *HeadscaleTransport) RoundTrip(req *http.Request) (*http.Response, error) {
-	recorder := httptest.NewRecorder()
-	t.verifyHandler(recorder, req)
-	resp := recorder.Result()
 	return resp, nil
 }

hscontrol/handlers.go

@@ -59,26 +59,35 @@ func parseCabailityVersion(req *http.Request) (tailcfg.CapabilityVersion, error)
 
 func (h *Headscale) handleVerifyRequest(
 	req *http.Request,
-) (bool, error) {
+	writer io.Writer,
+) error {
 	body, err := io.ReadAll(req.Body)
 	if err != nil {
-		return false, fmt.Errorf("cannot read request body: %w", err)
+		return fmt.Errorf("cannot read request body: %w", err)
 	}
 
 	var derpAdmitClientRequest tailcfg.DERPAdmitClientRequest
 	if err := json.Unmarshal(body, &derpAdmitClientRequest); err != nil {
-		return false, fmt.Errorf("cannot parse derpAdmitClientRequest: %w", err)
+		return fmt.Errorf("cannot parse derpAdmitClientRequest: %w", err)
 	}
 
 	nodes, err := h.db.ListNodes()
 	if err != nil {
-		return false, fmt.Errorf("cannot list nodes: %w", err)
+		return fmt.Errorf("cannot list nodes: %w", err)
 	}
 
-	return nodes.ContainsNodeKey(derpAdmitClientRequest.NodePublic), nil
+	resp := &tailcfg.DERPAdmitClientResponse{
+		Allow: nodes.ContainsNodeKey(derpAdmitClientRequest.NodePublic),
+	}
+	if err = json.NewEncoder(writer).Encode(resp); err != nil {
+		return fmt.Errorf("cannot encode response: %w", err)
 	}
 
-// see https://github.com/tailscale/tailscale/blob/964282d34f06ecc06ce644769c66b0b31d118340/derp/derp_server.go#L1159, Derp use verifyClientsURL to verify whether a client is allowed to connect to the DERP server.
+	return nil
+}
+
+// VerifyHandler see https://github.com/tailscale/tailscale/blob/964282d34f06ecc06ce644769c66b0b31d118340/derp/derp_server.go#L1159,
+// DERP use verifyClientsURL to verify whether a client is allowed to connect to the DERP server.
 func (h *Headscale) VerifyHandler(
 	writer http.ResponseWriter,
 	req *http.Request,
@@ -92,28 +101,18 @@ func (h *Headscale) VerifyHandler(
 		Str("handler", "/verify").
 		Msg("verify client")
 
-	allow, err := h.handleVerifyRequest(req)
+	if err := h.handleVerifyRequest(req, writer); err != nil {
-	if err != nil {
 		log.Error().
 			Caller().
 			Err(err).
 			Msg("Failed to verify client")
 		http.Error(writer, "Internal error", http.StatusInternalServerError)
-	}
 
-	resp := tailcfg.DERPAdmitClientResponse{
+		return
-		Allow: allow,
 	}
 
 	writer.Header().Set("Content-Type", "application/json")
 	writer.WriteHeader(http.StatusOK)
-	err = json.NewEncoder(writer).Encode(resp)
-	if err != nil {
-		log.Error().
-			Caller().
-			Err(err).
-			Msg("Failed to write response")
-	}
 }
 
 // KeyHandler provides the Headscale pub key