From 54c3e00a1ffb89d2b96d988f8fe21b3280a577f6 Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Sat, 5 Mar 2022 20:04:31 +0100
Subject: [PATCH] Merge local DERP server region with other configured DERP
 sources

---
 app.go         | 60 ++++++++------------------------------------------
 derp.go        |  1 +
 derp_server.go | 49 ++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 58 insertions(+), 52 deletions(-)

diff --git a/app.go b/app.go
index 4568558d..8cb987a9 100644
--- a/app.go
+++ b/app.go
@@ -13,7 +13,6 @@ import (
 	"os"
 	"os/signal"
 	"sort"
-	"strconv"
 	"strings"
 	"sync"
 	"syscall"
@@ -247,48 +246,6 @@ func NewHeadscale(cfg Config) (*Headscale, error) {
 			return nil, err
 		}
 		app.DERPServer = embeddedDERPServer
-
-		serverURL, err := url.Parse(app.cfg.ServerURL)
-		if err != nil {
-			return nil, err
-		}
-		var host string
-		var port int
-		host, portStr, err := net.SplitHostPort(serverURL.Host)
-		if err != nil {
-			if serverURL.Scheme == "https" {
-				host = serverURL.Host
-				port = 443
-			} else {
-				host = serverURL.Host
-				port = 80
-			}
-		} else {
-			port, err = strconv.Atoi(portStr)
-			if err != nil {
-				return nil, err
-			}
-		}
-
-		app.DERPMap = &tailcfg.DERPMap{
-			Regions: map[int]*tailcfg.DERPRegion{
-				999: {
-					RegionID:   999,
-					RegionCode: "headscale",
-					RegionName: "Headscale Embedded DERP",
-					Avoid:      false,
-					Nodes: []*tailcfg.DERPNode{
-						{
-							Name:     "999a",
-							RegionID: 999,
-							HostName: host,
-							DERPPort: port,
-						},
-					},
-				},
-			},
-			OmitDefaultRegions: false,
-		}
 	}
 
 	return &app, nil
@@ -536,17 +493,18 @@ func (h *Headscale) createRouter(grpcMux *runtime.ServeMux) *gin.Engine {
 func (h *Headscale) Serve() error {
 	var err error
 
+	// Fetch an initial DERP Map before we start serving
+	h.DERPMap = GetDERPMap(h.cfg.DERP)
+
 	if h.cfg.DERP.ServerEnabled {
 		go h.ServeSTUN()
-	} else {
-		// Fetch an initial DERP Map before we start serving
-		h.DERPMap = GetDERPMap(h.cfg.DERP)
+		h.DERPMap.Regions[h.DERPServer.region.RegionID] = &h.DERPServer.region
+	}
 
-		if h.cfg.DERP.AutoUpdate {
-			derpMapCancelChannel := make(chan struct{})
-			defer func() { derpMapCancelChannel <- struct{}{} }()
-			go h.scheduledDERPMapUpdateWorker(derpMapCancelChannel)
-		}
+	if h.cfg.DERP.AutoUpdate {
+		derpMapCancelChannel := make(chan struct{})
+		defer func() { derpMapCancelChannel <- struct{}{} }()
+		go h.scheduledDERPMapUpdateWorker(derpMapCancelChannel)
 	}
 
 	go h.expireEphemeralNodes(updateInterval)
diff --git a/derp.go b/derp.go
index 63e448db..7a9b2367 100644
--- a/derp.go
+++ b/derp.go
@@ -148,6 +148,7 @@ func (h *Headscale) scheduledDERPMapUpdateWorker(cancelChan <-chan struct{}) {
 		case <-ticker.C:
 			log.Info().Msg("Fetching DERPMap updates")
 			h.DERPMap = GetDERPMap(h.cfg.DERP)
+			h.DERPMap.Regions[h.DERPServer.region.RegionID] = &h.DERPServer.region
 
 			namespaces, err := h.ListNamespaces()
 			if err != nil {
diff --git a/derp_server.go b/derp_server.go
index e9009860..81cb1e29 100644
--- a/derp_server.go
+++ b/derp_server.go
@@ -6,6 +6,8 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"net/url"
+	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -14,6 +16,7 @@ import (
 	"github.com/rs/zerolog/log"
 	"tailscale.com/derp"
 	"tailscale.com/net/stun"
+	"tailscale.com/tailcfg"
 	"tailscale.com/types/key"
 )
 
@@ -30,12 +33,56 @@ var (
 
 type DERPServer struct {
 	tailscaleDERP *derp.Server
+	region        tailcfg.DERPRegion
 }
 
 func (h *Headscale) NewDERPServer() (*DERPServer, error) {
 	s := derp.NewServer(key.NodePrivate(*h.privateKey), log.Info().Msgf)
-	return &DERPServer{s}, nil
+	region, err := h.generateRegionLocalDERP()
+	if err != nil {
+		return nil, err
+	}
+	return &DERPServer{s, region}, nil
+}
 
+func (h *Headscale) generateRegionLocalDERP() (tailcfg.DERPRegion, error) {
+	serverURL, err := url.Parse(h.cfg.ServerURL)
+	if err != nil {
+		return tailcfg.DERPRegion{}, err
+	}
+	var host string
+	var port int
+	host, portStr, err := net.SplitHostPort(serverURL.Host)
+	if err != nil {
+		if serverURL.Scheme == "https" {
+			host = serverURL.Host
+			port = 443
+		} else {
+			host = serverURL.Host
+			port = 80
+		}
+	} else {
+		port, err = strconv.Atoi(portStr)
+		if err != nil {
+			return tailcfg.DERPRegion{}, err
+		}
+	}
+
+	localDERPregion := tailcfg.DERPRegion{
+		RegionID:   999,
+		RegionCode: "headscale",
+		RegionName: "Headscale Embedded DERP",
+		Avoid:      false,
+		Nodes: []*tailcfg.DERPNode{
+			{
+				Name:     "999a",
+				RegionID: 999,
+				HostName: host,
+				DERPPort: port,
+			},
+		},
+	}
+	return localDERPregion, nil
 }
 
 func (h *Headscale) DERPHandler(ctx *gin.Context) {