From ae0acf1084c2af9da23ae811e120097ad29e0a2d Mon Sep 17 00:00:00 2001
From: 117503445 <t117503445@gmail.com>
Date: Wed, 7 Aug 2024 23:40:41 +0800
Subject: [PATCH] feat: support client verify for derp

---
 hscontrol/app.go      |  2 ++
 hscontrol/handlers.go | 71 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/hscontrol/app.go b/hscontrol/app.go
index 5c85b064..37908283 100644
--- a/hscontrol/app.go
+++ b/hscontrol/app.go
@@ -457,6 +457,8 @@ func (h *Headscale) createRouter(grpcMux *grpcRuntime.ServeMux) *mux.Router {
 	router.HandleFunc("/swagger/v1/openapiv2.json", headscale.SwaggerAPIv1).
 		Methods(http.MethodGet)
 
+	router.HandleFunc("/verify", h.VerifyHandler).Methods(http.MethodPost)
+
 	if h.cfg.DERP.ServerEnabled {
 		router.HandleFunc("/derp", h.DERPServer.DERPHandler)
 		router.HandleFunc("/derp/probe", derpServer.DERPProbeHandler)
diff --git a/hscontrol/handlers.go b/hscontrol/handlers.go
index 72ec4e42..c97866d3 100644
--- a/hscontrol/handlers.go
+++ b/hscontrol/handlers.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"strconv"
 	"strings"
@@ -56,6 +57,76 @@ func parseCabailityVersion(req *http.Request) (tailcfg.CapabilityVersion, error)
 	return tailcfg.CapabilityVersion(clientCapabilityVersion), nil
 }
 
+// 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,
+) {
+	if req.Method != http.MethodPost {
+		http.Error(writer, "Wrong method", http.StatusMethodNotAllowed)
+		return
+	}
+	log.Debug().
+		Str("handler", "/verify").
+		Msg("verify client")
+
+	body, err := io.ReadAll(req.Body)
+	if err != nil {
+		log.Error().
+			Str("handler", "/verify").
+			Err(err).
+			Msg("Cannot read request body")
+		http.Error(writer, "Internal error", http.StatusInternalServerError)
+		return
+	}
+
+	var derpAdmitClientRequest tailcfg.DERPAdmitClientRequest
+	if err := json.Unmarshal(body, &derpAdmitClientRequest); err != nil {
+		log.Error().
+			Caller().
+			Err(err).
+			Msg("Cannot parse derpAdmitClientRequest")
+		http.Error(writer, "Internal error", http.StatusInternalServerError)
+		return
+	}
+
+	nodes, err := h.db.ListNodes()
+	if err != nil {
+		log.Error().
+			Caller().
+			Err(err).
+			Msg("Cannot list nodes")
+		http.Error(writer, "Internal error", http.StatusInternalServerError)
+	}
+
+	for _, node := range nodes {
+		log.Debug().Str("node", node.NodeKey.String()).Msg("Node")
+	}
+
+	allow := false
+	// Check if the node is in the list of nodes
+	for _, node := range nodes {
+		if node.NodeKey == derpAdmitClientRequest.NodePublic {
+			allow = true
+			break
+		}
+	}
+
+	resp := tailcfg.DERPAdmitClientResponse{
+		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
 // Listens in /key.
 func (h *Headscale) KeyHandler(