From afdae0742874a7cc16f62d9e52dc0c537835cbf9 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Sat, 18 Nov 2023 19:09:44 +0100
Subject: [PATCH] use key type in notifier

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>
---
 hscontrol/notifier/notifier.go | 21 +++++++++++----------
 hscontrol/poll.go              |  4 ++--
 2 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/hscontrol/notifier/notifier.go b/hscontrol/notifier/notifier.go
index 32f426ad72..fee0befb59 100644
--- a/hscontrol/notifier/notifier.go
+++ b/hscontrol/notifier/notifier.go
@@ -6,6 +6,7 @@ import (
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/juanfont/headscale/hscontrol/util"
 	"github.com/rs/zerolog/log"
+	"tailscale.com/types/key"
 )
 
 type Notifier struct {
@@ -17,9 +18,9 @@ func NewNotifier() *Notifier {
 	return &Notifier{}
 }
 
-func (n *Notifier) AddNode(machineKey string, c chan<- types.StateUpdate) {
-	log.Trace().Caller().Str("key", machineKey).Msg("acquiring lock to add node")
-	defer log.Trace().Caller().Str("key", machineKey).Msg("releasing lock to add node")
+func (n *Notifier) AddNode(machineKey key.MachinePublic, c chan<- types.StateUpdate) {
+	log.Trace().Caller().Str("key", machineKey.ShortString()).Msg("acquiring lock to add node")
+	defer log.Trace().Caller().Str("key", machineKey.ShortString()).Msg("releasing lock to add node")
 
 	n.l.Lock()
 	defer n.l.Unlock()
@@ -28,17 +29,17 @@ func (n *Notifier) AddNode(machineKey string, c chan<- types.StateUpdate) {
 		n.nodes = make(map[string]chan<- types.StateUpdate)
 	}
 
-	n.nodes[machineKey] = c
+	n.nodes[machineKey.String()] = c
 
 	log.Trace().
-		Str("machine_key", machineKey).
+		Str("machine_key", machineKey.ShortString()).
 		Int("open_chans", len(n.nodes)).
 		Msg("Added new channel")
 }
 
-func (n *Notifier) RemoveNode(machineKey string) {
-	log.Trace().Caller().Str("key", machineKey).Msg("acquiring lock to remove node")
-	defer log.Trace().Caller().Str("key", machineKey).Msg("releasing lock to remove node")
+func (n *Notifier) RemoveNode(machineKey key.MachinePublic) {
+	log.Trace().Caller().Str("key", machineKey.ShortString()).Msg("acquiring lock to remove node")
+	defer log.Trace().Caller().Str("key", machineKey.ShortString()).Msg("releasing lock to remove node")
 
 	n.l.Lock()
 	defer n.l.Unlock()
@@ -47,10 +48,10 @@ func (n *Notifier) RemoveNode(machineKey string) {
 		return
 	}
 
-	delete(n.nodes, machineKey)
+	delete(n.nodes, machineKey.String())
 
 	log.Trace().
-		Str("machine_key", machineKey).
+		Str("machine_key", machineKey.ShortString()).
 		Int("open_chans", len(n.nodes)).
 		Msg("Removed channel")
 }
diff --git a/hscontrol/poll.go b/hscontrol/poll.go
index 4713a2e898..050e857a40 100644
--- a/hscontrol/poll.go
+++ b/hscontrol/poll.go
@@ -232,8 +232,8 @@ func (h *Headscale) handlePoll(
 	defer closeChanWithLog(updateChan, node.Hostname, "updateChan")
 
 	// Register the node's update channel
-	h.nodeNotifier.AddNode(node.MachineKey.String(), updateChan)
-	defer h.nodeNotifier.RemoveNode(node.MachineKey.String())
+	h.nodeNotifier.AddNode(node.MachineKey, updateChan)
+	defer h.nodeNotifier.RemoveNode(node.MachineKey)
 
 	keepAliveTicker := time.NewTicker(keepAliveInterval)