From 204dedaa4901f43991561d05650c8ff6ab6b427c Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Sun, 4 Sep 2022 15:14:12 +0200
Subject: [PATCH] Only pass the context in pollmap, no req needed

---
 protocol_common_poll.go | 8 ++++----
 protocol_legacy_poll.go | 2 +-
 protocol_noise_poll.go  | 2 +-
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/protocol_common_poll.go b/protocol_common_poll.go
index 65dcb55..6dedfd0 100644
--- a/protocol_common_poll.go
+++ b/protocol_common_poll.go
@@ -22,7 +22,7 @@ const machineNameContextKey = contextKey("machineName")
 // managed the poll loop.
 func (h *Headscale) handlePollCommon(
 	writer http.ResponseWriter,
-	req *http.Request,
+	ctx context.Context,
 	machine *Machine,
 	mapRequest tailcfg.MapRequest,
 	isNoise bool,
@@ -201,7 +201,7 @@ func (h *Headscale) handlePollCommon(
 
 	h.pollNetMapStream(
 		writer,
-		req,
+		ctx,
 		machine,
 		mapRequest,
 		pollDataChan,
@@ -221,7 +221,7 @@ func (h *Headscale) handlePollCommon(
 // ensuring we communicate updates and data to the connected clients.
 func (h *Headscale) pollNetMapStream(
 	writer http.ResponseWriter,
-	req *http.Request,
+	ctxReq context.Context,
 	machine *Machine,
 	mapRequest tailcfg.MapRequest,
 	pollDataChan chan []byte,
@@ -232,7 +232,7 @@ func (h *Headscale) pollNetMapStream(
 	h.pollNetMapStreamWG.Add(1)
 	defer h.pollNetMapStreamWG.Done()
 
-	ctx := context.WithValue(req.Context(), machineNameContextKey, machine.Hostname)
+	ctx := context.WithValue(ctxReq, machineNameContextKey, machine.Hostname)
 
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
diff --git a/protocol_legacy_poll.go b/protocol_legacy_poll.go
index f7ef654..f27ee4e 100644
--- a/protocol_legacy_poll.go
+++ b/protocol_legacy_poll.go
@@ -90,5 +90,5 @@ func (h *Headscale) PollNetMapHandler(
 		Str("machine", machine.Hostname).
 		Msg("A machine is entering polling via the legacy protocol")
 
-	h.handlePollCommon(writer, req, machine, mapRequest, false)
+	h.handlePollCommon(writer, req.Context(), machine, mapRequest, false)
 }
diff --git a/protocol_noise_poll.go b/protocol_noise_poll.go
index 8498dcf..b15183c 100644
--- a/protocol_noise_poll.go
+++ b/protocol_noise_poll.go
@@ -63,5 +63,5 @@ func (h *Headscale) NoisePollNetMapHandler(
 		Str("machine", machine.Hostname).
 		Msg("A machine is entering polling via the Noise protocol")
 
-	h.handlePollCommon(writer, req, machine, mapRequest, true)
+	h.handlePollCommon(writer, req.Context(), machine, mapRequest, true)
 }