Merge pull request #278 from enoperm/pollnetmap-update-only

juanfont · Jan 29, 2022 · b122d06 · b122d06
2 parents 052fccd + 6f9ed95
commit b122d06
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 17 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,9 @@
 
 **TBD (TBD):**
 
+- Fixed issue where hosts deleted from control server may be written back to the database, as long as they are connected to the control server [#278](https://github.com/juanfont/headscale/pull/278)
+  )
+
 **0.12.3 (2022-01-13):**
 
 **Changes**:

diff --git 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
@@ -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,14 +443,26 @@ 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()
 
 				lastStateUpdate.WithLabelValues(machine.Namespace.Name, machine.Name).
 					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").