From 68dc2a70dbe789a7b7de78439dd52d0a6f0a0cab Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Mon, 22 Nov 2021 19:51:16 +0000
Subject: [PATCH] Update neighbours if node is expired or refreshed

In addition, only pass the map of registered and not expired nodes to
clients.
---
 api.go     |  2 +-
 machine.go | 22 ++++++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/api.go b/api.go
index 5a0b9d37..6b75ad45 100644
--- a/api.go
+++ b/api.go
@@ -186,7 +186,7 @@ func (h *Headscale) getMapResponse(
 		return nil, err
 	}
 
-	peers, err := h.getPeers(machine)
+	peers, err := h.getValidPeers(machine)
 	if err != nil {
 		log.Error().
 			Str("func", "getMapResponse").
diff --git a/machine.go b/machine.go
index 06fec743..306b3a47 100644
--- a/machine.go
+++ b/machine.go
@@ -207,6 +207,23 @@ func (h *Headscale) getPeers(machine *Machine) (Machines, error) {
 	return peers, nil
 }
 
+func (h *Headscale) getValidPeers(machine *Machine) (Machines, error) {
+	validPeers := make(Machines, 0)
+
+	peers, err := h.getPeers(machine)
+	if err != nil {
+		return Machines{}, err
+	}
+
+	for _, peer := range peers {
+		if peer.isRegistered() && !peer.isExpired() {
+			validPeers = append(validPeers, peer)
+		}
+	}
+
+	return validPeers, nil
+}
+
 func (h *Headscale) ListMachines() ([]Machine, error) {
 	machines := []Machine{}
 	if err := h.db.Preload("AuthKey").Preload("AuthKey.Namespace").Preload("Namespace").Find(&machines).Error; err != nil {
@@ -267,6 +284,8 @@ func (h *Headscale) ExpireMachine(machine *Machine) {
 	now := time.Now()
 	machine.Expiry = &now
 
+	h.setLastStateChangeToNow(machine.Namespace.Name)
+
 	h.db.Save(machine)
 }
 
@@ -276,6 +295,9 @@ func (h *Headscale) RefreshMachine(machine *Machine, expiry time.Time) {
 
 	machine.LastSuccessfulUpdate = &now
 	machine.Expiry = &expiry
+
+	h.setLastStateChangeToNow(machine.Namespace.Name)
+
 	h.db.Save(machine)
 }