From a87a9636e3f1432b3f944fc7dd81a7adba198d75 Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Fri, 19 Aug 2022 14:19:29 +0200
Subject: [PATCH] Expanded response marshal methods to support legacy and Noise

---
 protocol_common_utils.go | 29 ++++++++++++++++++++++++-----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/protocol_common_utils.go b/protocol_common_utils.go
index 516b616c..3dc435f3 100644
--- a/protocol_common_utils.go
+++ b/protocol_common_utils.go
@@ -21,7 +21,7 @@ func (h *Headscale) getMapResponseData(
 	}
 
 	if isNoise {
-		return h.marshalResponse(mapResponse, mapRequest.Compress, key.MachinePublic{})
+		return h.marshalMapResponse(mapResponse, key.MachinePublic{}, mapRequest.Compress)
 	}
 
 	var machineKey key.MachinePublic
@@ -35,7 +35,7 @@ func (h *Headscale) getMapResponseData(
 		return nil, err
 	}
 
-	return h.marshalResponse(mapResponse, mapRequest.Compress, machineKey)
+	return h.marshalMapResponse(mapResponse, machineKey, mapRequest.Compress)
 }
 
 func (h *Headscale) getMapKeepAliveResponseData(
@@ -48,7 +48,7 @@ func (h *Headscale) getMapKeepAliveResponseData(
 	}
 
 	if isNoise {
-		return h.marshalResponse(keepAliveResponse, mapRequest.Compress, key.MachinePublic{})
+		return h.marshalMapResponse(keepAliveResponse, key.MachinePublic{}, mapRequest.Compress)
 	}
 
 	var machineKey key.MachinePublic
@@ -62,13 +62,32 @@ func (h *Headscale) getMapKeepAliveResponseData(
 		return nil, err
 	}
 
-	return h.marshalResponse(keepAliveResponse, mapRequest.Compress, machineKey)
+	return h.marshalMapResponse(keepAliveResponse, machineKey, mapRequest.Compress)
 }
 
 func (h *Headscale) marshalResponse(
 	resp interface{},
-	compression string,
 	machineKey key.MachinePublic,
+) ([]byte, error) {
+	jsonBody, err := json.Marshal(resp)
+	if err != nil {
+		log.Error().
+			Caller().
+			Err(err).
+			Msg("Cannot marshal response")
+	}
+
+	if machineKey.IsZero() { // if Noise
+		return jsonBody, nil
+	}
+
+	return h.privateKey.SealTo(machineKey, jsonBody), nil
+}
+
+func (h *Headscale) marshalMapResponse(
+	resp interface{},
+	machineKey key.MachinePublic,
+	compression string,
 ) ([]byte, error) {
 	jsonBody, err := json.Marshal(resp)
 	if err != nil {