From 39b85b02bbef9f0ecd74112524e883d3b3dcbd59 Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Sun, 14 Aug 2022 03:20:53 +0200
Subject: [PATCH] Move getMapResponse into reusable function by TS2019 and
 TS2021

---
 api.go  | 140 ++++++++++++++++++++++++++++++--------------------------
 poll.go |   4 +-
 2 files changed, 78 insertions(+), 66 deletions(-)

diff --git a/api.go b/api.go
index 82e13c77..86d77b6a 100644
--- a/api.go
+++ b/api.go
@@ -396,78 +396,16 @@ func (h *Headscale) RegistrationHandler(
 	}
 }
 
-func (h *Headscale) getMapResponse(
+func (h *Headscale) getLegacyMapResponseData(
 	machineKey key.MachinePublic,
 	mapRequest tailcfg.MapRequest,
 	machine *Machine,
 ) ([]byte, error) {
-	log.Trace().
-		Str("func", "getMapResponse").
-		Str("machine", mapRequest.Hostinfo.Hostname).
-		Msg("Creating Map response")
-	node, err := machine.toNode(h.cfg.BaseDomain, h.cfg.DNSConfig, true)
+	resp, err := h.generateMapResponse(mapRequest, machine)
 	if err != nil {
-		log.Error().
-			Caller().
-			Str("func", "getMapResponse").
-			Err(err).
-			Msg("Cannot convert to node")
-
 		return nil, err
 	}
 
-	peers, err := h.getValidPeers(machine)
-	if err != nil {
-		log.Error().
-			Caller().
-			Str("func", "getMapResponse").
-			Err(err).
-			Msg("Cannot fetch peers")
-
-		return nil, err
-	}
-
-	profiles := getMapResponseUserProfiles(*machine, peers)
-
-	nodePeers, err := peers.toNodes(h.cfg.BaseDomain, h.cfg.DNSConfig, true)
-	if err != nil {
-		log.Error().
-			Caller().
-			Str("func", "getMapResponse").
-			Err(err).
-			Msg("Failed to convert peers to Tailscale nodes")
-
-		return nil, err
-	}
-
-	dnsConfig := getMapResponseDNSConfig(
-		h.cfg.DNSConfig,
-		h.cfg.BaseDomain,
-		*machine,
-		peers,
-	)
-
-	resp := tailcfg.MapResponse{
-		KeepAlive:    false,
-		Node:         node,
-		Peers:        nodePeers,
-		DNSConfig:    dnsConfig,
-		Domain:       h.cfg.BaseDomain,
-		PacketFilter: h.aclRules,
-		DERPMap:      h.DERPMap,
-		UserProfiles: profiles,
-		Debug: &tailcfg.Debug{
-			DisableLogTail:      !h.cfg.LogTail.Enabled,
-			RandomizeClientPort: h.cfg.RandomizeClientPort,
-		},
-	}
-
-	log.Trace().
-		Str("func", "getMapResponse").
-		Str("machine", mapRequest.Hostinfo.Hostname).
-		// Interface("payload", resp).
-		Msgf("Generated map response: %s", tailMapResponseToString(resp))
-
 	var respBody []byte
 	if mapRequest.Compress == "zstd" {
 		src, err := json.Marshal(resp)
@@ -498,6 +436,80 @@ func (h *Headscale) getMapResponse(
 	return data, nil
 }
 
+func (h *Headscale) generateMapResponse(
+	mapRequest tailcfg.MapRequest,
+	machine *Machine,
+) (*tailcfg.MapResponse, error) {
+	log.Trace().
+		Str("func", "generateMapResponse").
+		Str("machine", mapRequest.Hostinfo.Hostname).
+		Msg("Creating Map response")
+	node, err := machine.toNode(h.cfg.BaseDomain, h.cfg.DNSConfig, true)
+	if err != nil {
+		log.Error().
+			Caller().
+			Str("func", "generateMapResponse").
+			Err(err).
+			Msg("Cannot convert to node")
+
+		return nil, err
+	}
+
+	peers, err := h.getValidPeers(machine)
+	if err != nil {
+		log.Error().
+			Caller().
+			Str("func", "generateMapResponse").
+			Err(err).
+			Msg("Cannot fetch peers")
+
+		return nil, err
+	}
+
+	profiles := getMapResponseUserProfiles(*machine, peers)
+
+	nodePeers, err := peers.toNodes(h.cfg.BaseDomain, h.cfg.DNSConfig, true)
+	if err != nil {
+		log.Error().
+			Caller().
+			Str("func", "generateMapResponse").
+			Err(err).
+			Msg("Failed to convert peers to Tailscale nodes")
+
+		return nil, err
+	}
+
+	dnsConfig := getMapResponseDNSConfig(
+		h.cfg.DNSConfig,
+		h.cfg.BaseDomain,
+		*machine,
+		peers,
+	)
+
+	resp := tailcfg.MapResponse{
+		KeepAlive:    false,
+		Node:         node,
+		Peers:        nodePeers,
+		DNSConfig:    dnsConfig,
+		Domain:       h.cfg.BaseDomain,
+		PacketFilter: h.aclRules,
+		DERPMap:      h.DERPMap,
+		UserProfiles: profiles,
+		Debug: &tailcfg.Debug{
+			DisableLogTail:      !h.cfg.LogTail.Enabled,
+			RandomizeClientPort: h.cfg.RandomizeClientPort,
+		},
+	}
+
+	log.Trace().
+		Str("func", "generateMapResponse").
+		Str("machine", mapRequest.Hostinfo.Hostname).
+		// Interface("payload", resp).
+		Msgf("Generated map response: %s", tailMapResponseToString(resp))
+
+	return &resp, nil
+}
+
 func (h *Headscale) getMapKeepAliveResponse(
 	machineKey key.MachinePublic,
 	mapRequest tailcfg.MapRequest,
diff --git a/poll.go b/poll.go
index 9c17b5cb..a51c9362 100644
--- a/poll.go
+++ b/poll.go
@@ -143,7 +143,7 @@ func (h *Headscale) PollNetMapHandler(
 		}
 	}
 
-	data, err := h.getMapResponse(machineKey, mapRequest, machine)
+	data, err := h.getLegacyMapResponseData(machineKey, mapRequest, machine)
 	if err != nil {
 		log.Error().
 			Str("handler", "PollNetMap").
@@ -491,7 +491,7 @@ func (h *Headscale) PollNetMapStream(
 					Time("last_successful_update", lastUpdate).
 					Time("last_state_change", h.getLastStateChange(machine.Namespace.Name)).
 					Msgf("There has been updates since the last successful update to %s", machine.Hostname)
-				data, err := h.getMapResponse(machineKey, mapRequest, machine)
+				data, err := h.getLegacyMapResponseData(machineKey, mapRequest, machine)
 				if err != nil {
 					log.Error().
 						Str("handler", "PollNetMapStream").