From 25ccf5ef1808375a17337e9f1a2562858c5edb7e Mon Sep 17 00:00:00 2001
From: Csaba Sarkadi <sarkadicsa@tutanota.de>
Date: Sun, 16 Jan 2022 11:59:03 +0100
Subject: [PATCH] PollNetMapStream: do not create any rows during long-poll
 operation

---
 machine.go |  8 ++++++
 poll.go    | 84 +++++++++++++++++++++++++++++++++++++++++++-----------
 2 files changed, 75 insertions(+), 17 deletions(-)

diff --git a/machine.go b/machine.go
index d58c9c5..de1764e 100644
--- a/machine.go
+++ b/machine.go
@@ -319,6 +319,14 @@ func (h *Headscale) DeleteMachine(machine *Machine) error {
 	return h.RequestMapUpdates(namespaceID)
 }
 
+func (h *Headscale) TouchMachine(machine *Machine) error {
+	return h.db.Updates(Machine{
+		ID:                   machine.ID,
+		LastSeen:             machine.LastSeen,
+		LastSuccessfulUpdate: machine.LastSuccessfulUpdate,
+	}).Error
+}
+
 // HardDeleteMachine hard deletes a Machine from the database.
 func (h *Headscale) HardDeleteMachine(machine *Machine) error {
 	err := h.RemoveSharedMachineFromAllNamespaces(machine)
diff --git a/poll.go b/poll.go
index 1038903..d7aa12e 100644
--- a/poll.go
+++ b/poll.go
@@ -102,7 +102,7 @@ func (h *Headscale) PollNetMapHandler(ctx *gin.Context) {
 		machine.Endpoints = datatypes.JSON(endpoints)
 		machine.LastSeen = &now
 	}
-	h.db.Save(&machine)
+	h.db.Updates(machine)
 
 	data, err := h.getMapResponse(machineKey, req, machine)
 	if err != nil {
@@ -291,6 +291,10 @@ func (h *Headscale) PollNetMapStream(
 					Str("channel", "pollData").
 					Err(err).
 					Msg("Cannot update machine from database")
+
+				// client has been removed from database
+				// since the stream opened, terminate connection.
+				return false
 			}
 			now := time.Now().UTC()
 			machine.LastSeen = &now
@@ -299,13 +303,22 @@ func (h *Headscale) PollNetMapStream(
 				Set(float64(now.Unix()))
 			machine.LastSuccessfulUpdate = &now
 
-			h.db.Save(&machine)
-			log.Trace().
-				Str("handler", "PollNetMapStream").
-				Str("machine", machine.Name).
-				Str("channel", "pollData").
-				Int("bytes", len(data)).
-				Msg("Machine entry in database updated successfully after sending pollData")
+			err = h.TouchMachine(machine)
+			if err != nil {
+				log.Error().
+					Str("handler", "PollNetMapStream").
+					Str("machine", machine.Name).
+					Str("channel", "pollData").
+					Err(err).
+					Msg("Cannot update machine LastSuccessfulUpdate")
+			} else {
+				log.Trace().
+					Str("handler", "PollNetMapStream").
+					Str("machine", machine.Name).
+					Str("channel", "pollData").
+					Int("bytes", len(data)).
+					Msg("Machine entry in database updated successfully after sending pollData")
+			}
 
 			return true
 
@@ -344,16 +357,29 @@ func (h *Headscale) PollNetMapStream(
 					Str("channel", "keepAlive").
 					Err(err).
 					Msg("Cannot update machine from database")
+
+				// client has been removed from database
+				// since the stream opened, terminate connection.
+				return false
 			}
 			now := time.Now().UTC()
 			machine.LastSeen = &now
-			h.db.Save(&machine)
-			log.Trace().
-				Str("handler", "PollNetMapStream").
-				Str("machine", machine.Name).
-				Str("channel", "keepAlive").
-				Int("bytes", len(data)).
-				Msg("Machine updated successfully after sending keep alive")
+			err = h.TouchMachine(machine)
+			if err != nil {
+				log.Error().
+					Str("handler", "PollNetMapStream").
+					Str("machine", machine.Name).
+					Str("channel", "keepAlive").
+					Err(err).
+					Msg("Cannot update machine LastSeen")
+			} else {
+				log.Trace().
+					Str("handler", "PollNetMapStream").
+					Str("machine", machine.Name).
+					Str("channel", "keepAlive").
+					Int("bytes", len(data)).
+					Msg("Machine updated successfully after sending keep alive")
+			}
 
 			return true
 
@@ -417,6 +443,10 @@ func (h *Headscale) PollNetMapStream(
 						Str("channel", "update").
 						Err(err).
 						Msg("Cannot update machine from database")
+
+					// client has been removed from database
+					// since the stream opened, terminate connection.
+					return false
 				}
 				now := time.Now().UTC()
 
@@ -424,7 +454,15 @@ func (h *Headscale) PollNetMapStream(
 					Set(float64(now.Unix()))
 				machine.LastSuccessfulUpdate = &now
 
-				h.db.Save(&machine)
+				err = h.TouchMachine(machine)
+				if err != nil {
+					log.Error().
+						Str("handler", "PollNetMapStream").
+						Str("machine", machine.Name).
+						Str("channel", "update").
+						Err(err).
+						Msg("Cannot update machine LastSuccessfulUpdate")
+				}
 			} else {
 				log.Trace().
 					Str("handler", "PollNetMapStream").
@@ -452,10 +490,22 @@ func (h *Headscale) PollNetMapStream(
 					Str("channel", "Done").
 					Err(err).
 					Msg("Cannot update machine from database")
+
+				// client has been removed from database
+				// since the stream opened, terminate connection.
+				return false
 			}
 			now := time.Now().UTC()
 			machine.LastSeen = &now
-			h.db.Save(&machine)
+			err = h.TouchMachine(machine)
+			if err != nil {
+				log.Error().
+					Str("handler", "PollNetMapStream").
+					Str("machine", machine.Name).
+					Str("channel", "Done").
+					Err(err).
+					Msg("Cannot update machine LastSeen")
+			}
 
 			log.Trace().
 				Str("handler", "PollNetMapStream").