From eb500155e84847774d79b8454c0327e37c23c0ff Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Sun, 6 Mar 2022 17:00:56 +0100
Subject: [PATCH] Make STUN server configurable

---
 app.go                     |  6 +++++-
 cmd/headscale/cli/utils.go |  4 ++++
 config-example.yaml        |  6 ++++++
 derp_server.go             | 22 ++++++++++++++++++----
 4 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/app.go b/app.go
index 8cb987a9..82e87cf0 100644
--- a/app.go
+++ b/app.go
@@ -121,6 +121,8 @@ type OIDCConfig struct {
 
 type DERPConfig struct {
 	ServerEnabled   bool
+	STUNEnabled     bool
+	STUNAddr        string
 	URLs            []url.URL
 	Paths           []string
 	AutoUpdate      bool
@@ -497,8 +499,10 @@ func (h *Headscale) Serve() error {
 	h.DERPMap = GetDERPMap(h.cfg.DERP)
 
 	if h.cfg.DERP.ServerEnabled {
-		go h.ServeSTUN()
 		h.DERPMap.Regions[h.DERPServer.region.RegionID] = &h.DERPServer.region
+		if h.cfg.DERP.STUNEnabled {
+			go h.ServeSTUN()
+		}
 	}
 
 	if h.cfg.DERP.AutoUpdate {
diff --git a/cmd/headscale/cli/utils.go b/cmd/headscale/cli/utils.go
index 7277723f..e6dce3a1 100644
--- a/cmd/headscale/cli/utils.go
+++ b/cmd/headscale/cli/utils.go
@@ -118,6 +118,8 @@ func LoadConfig(path string) error {
 
 func GetDERPConfig() headscale.DERPConfig {
 	enabled := viper.GetBool("derp.server.enabled")
+	stunEnabled := viper.GetBool("derp.server.stun.enabled")
+	stunAddr := viper.GetString("derp.server.stun.listen_addr")
 
 	urlStrs := viper.GetStringSlice("derp.urls")
 
@@ -141,6 +143,8 @@ func GetDERPConfig() headscale.DERPConfig {
 
 	return headscale.DERPConfig{
 		ServerEnabled:   enabled,
+		STUNEnabled:     stunEnabled,
+		STUNAddr:        stunAddr,
 		URLs:            urls,
 		Paths:           paths,
 		AutoUpdate:      autoUpdate,
diff --git a/config-example.yaml b/config-example.yaml
index 6f4060ac..57b43fd4 100644
--- a/config-example.yaml
+++ b/config-example.yaml
@@ -60,6 +60,12 @@ derp:
     # The Headscale server_url defined above MUST be using https, DERP requires TLS to be in place
     enabled: false
 
+    # If enabled, also listens in the configured address for STUN connections to help on NAT traversal
+    # For more details on how this works, check this great article: https://tailscale.com/blog/how-tailscale-works/
+    stun:
+      enabled: false
+      listen_addr: "0.0.0.0:3478"
+
   # List of externally available DERP maps encoded in JSON
   urls:
     - https://controlplane.tailscale.com/derpmap/default
diff --git a/derp_server.go b/derp_server.go
index dcda63e9..aeb4877a 100644
--- a/derp_server.go
+++ b/derp_server.go
@@ -75,6 +75,19 @@ func (h *Headscale) generateRegionLocalDERP() (tailcfg.DERPRegion, error) {
 			},
 		},
 	}
+
+	if h.cfg.DERP.STUNEnabled {
+		_, portStr, err := net.SplitHostPort(h.cfg.DERP.STUNAddr)
+		if err != nil {
+			return tailcfg.DERPRegion{}, err
+		}
+		port, err := strconv.Atoi(portStr)
+		if err != nil {
+			return tailcfg.DERPRegion{}, err
+		}
+		localDERPregion.Nodes[0].STUNPort = port
+	}
+
 	return localDERPregion, nil
 }
 
@@ -136,6 +149,7 @@ func (h *Headscale) DERPProbeHandler(ctx *gin.Context) {
 // because its DNS are broken.
 // The initial implementation is here https://github.com/tailscale/tailscale/pull/1406
 // They have a cache, but not clear if that is really necessary at Headscale, uh, scale.
+// An example implementation is found here https://derp.tailscale.com/bootstrap-dns
 func (h *Headscale) DERPBootstrapDNSHandler(ctx *gin.Context) {
 	dnsEntries := make(map[string][]net.IP)
 
@@ -155,14 +169,14 @@ func (h *Headscale) DERPBootstrapDNSHandler(ctx *gin.Context) {
 	ctx.JSON(http.StatusOK, dnsEntries)
 }
 
-// ServeSTUN starts a STUN server on udp/3478
+// ServeSTUN starts a STUN server on the configured addr
 func (h *Headscale) ServeSTUN() {
-	pc, err := net.ListenPacket("udp", "0.0.0.0:3478")
+	packetConn, err := net.ListenPacket("udp", h.cfg.DERP.STUNAddr)
 	if err != nil {
 		log.Fatal().Msgf("failed to open STUN listener: %v", err)
 	}
-	log.Trace().Msgf("STUN server started at %s", pc.LocalAddr())
-	serverSTUNListener(context.Background(), pc.(*net.UDPConn))
+	log.Info().Msgf("STUN server started at %s", packetConn.LocalAddr())
+	serverSTUNListener(context.Background(), packetConn.(*net.UDPConn))
 }
 
 func serverSTUNListener(ctx context.Context, pc *net.UDPConn) {