update peers immediatly if a client disconnects

Signed-off-by: Kristoffer Dalby <[email protected]>

hscontrol/mapper/mapper.go

@@ -335,6 +335,41 @@ func (m *Mapper) PeerChangedResponse(
 	return m.marshalMapResponse(mapRequest, &resp, node, mapRequest.Compress)
 }
 
+func (m *Mapper) PeerOfflineChangedResponse(
+	mapRequest tailcfg.MapRequest,
+	node *types.Node,
+	changed types.Nodes,
+	pol *policy.ACLPolicy,
+) ([]byte, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	offlineMap := map[tailcfg.NodeID]bool{}
+
+	// Update our internal map.
+	for _, node := range changed {
+		offlineMap[tailcfg.NodeID(node.ID)] = false
+	}
+
+	resp := m.baseMapResponse()
+
+	tailnode, err := tailNode(node, m.capVer, pol, m.dnsCfg, m.baseDomain, m.randomClientPort)
+	if err != nil {
+		return nil, err
+	}
+	resp.Node = tailnode
+
+	// PeerSeenChange contains information on how to update peers' LastSeen
+	// times. If the value is false, the peer is gone. If the value is true,
+	// the LastSeen time is now. Absent means unchanged.
+	resp.PeerSeenChange = offlineMap
+
+	// OnlineChange changes the value of a Peer Node.Online value.
+	resp.OnlineChange = offlineMap
+
+	return m.marshalMapResponse(mapRequest, &resp, node, mapRequest.Compress)
+}
+
 // PeerOnlineChanged is a no response function that internally updates
 // the mappers internal node map with LastSeen and Online information.
 func (m *Mapper) PeerOnlineChanged(

hscontrol/poll.go

@@ -243,6 +243,15 @@ func (h *Headscale) handlePoll(
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
+	// If the streaming session is closed, let the other nodes know
+	// that this node has gone away.
+	defer func() {
+		h.nodeNotifier.NotifyWithIgnore(types.StateUpdate{
+			Type:    types.StatePeerOfflineChanged,
+			Changed: types.Nodes{node},
+		}, node.MachineKey)
+	}()
+
 	for {
 		logInfo("Waiting for update on stream channel")
 		select {
@@ -285,6 +294,9 @@ func (h *Headscale) handlePoll(
 			case types.StatePeerChanged:
 				logInfo("Sending PeerChanged MapResponse")
 				data, err = mapp.PeerChangedResponse(mapRequest, node, update.Changed, h.ACLPolicy)
+			case types.StatePeerOfflineChanged:
+				logInfo("Sending PeerOffline MapResponse")
+				data, err = mapp.PeerOfflineChangedResponse(mapRequest, node, update.Changed, h.ACLPolicy)
 			case types.StatePeerRemoved:
 				logInfo("Sending PeerRemoved MapResponse")
 				data, err = mapp.PeerRemovedResponse(mapRequest, node, update.Removed)

hscontrol/types/common.go

@@ -115,6 +115,7 @@ const (
 	StatePeerRemoved
 	StateDERPUpdated
 	StatePeerOnlineChanged
+	StatePeerOfflineChanged
 )
 
 // StateUpdate is an internal message containing information about