MyFavMirrors/headscale
1
0
Fork 0
mirror of https://github.com/juanfont/headscale.git synced 2024-12-25 05:35:52 -05:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #34 from juanfont/fix-polling-race

Browse Source 
Fix a race condition when a client closes the connection
This commit is contained in:
Juan Font 2021-06-09 20:56:58 +02:00 committed by GitHub
commit 6e86b2af0c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
api.go
View File

@ -241,7 +241,6 @@ func (h *Headscale) PollNetMapHandler(c *gin.Context) {
update := make(chan []byte, 1) update := make(chan []byte, 1)
cancelKeepAlive := make(chan []byte, 1) cancelKeepAlive := make(chan []byte, 1)
defer close(pollData) defer close(pollData)
defer close(update)
defer close(cancelKeepAlive) defer close(cancelKeepAlive)
h.pollMu.Lock() h.pollMu.Lock()
h.clientsPolling[m.ID] = update h.clientsPolling[m.ID] = update
@ -283,8 +282,9 @@ func (h *Headscale) PollNetMapHandler(c *gin.Context) {
peers, _ := h.getPeers(m) peers, _ := h.getPeers(m)
h.pollMu.Lock() h.pollMu.Lock()
for _, p := range *peers { for _, p := range *peers {
log.Printf("[%s] Notifying peer %s (%s)", m.Name, p.Name, p.Addresses[0]) pUp, ok := h.clientsPolling[uint64(p.ID)]
if pUp, ok := h.clientsPolling[uint64(p.ID)]; ok { if ok {
log.Printf("[%s] Notifying peer %s (%s)", m.Name, p.Name, p.Addresses[0])
pUp <- []byte{} pUp <- []byte{}
} else { } else {
log.Printf("[%s] Peer %s does not appear to be polling", m.Name, p.Name) log.Printf("[%s] Peer %s does not appear to be polling", m.Name, p.Name)
@ -311,22 +311,23 @@ func (h *Headscale) PollNetMapHandler(c *gin.Context) {
log.Printf("[%s] Received a request for update", m.Name) log.Printf("[%s] Received a request for update", m.Name)
data, err := h.getMapResponse(mKey, req, m) data, err := h.getMapResponse(mKey, req, m)
if err != nil { if err != nil {
log.Printf("[%s] 🤮 Cannot get the poll response: %s", m.Name, err) log.Printf("[%s] Could not get the map update: %s", m.Name, err)
} }
_, err = w.Write(*data) _, err = w.Write(*data)
if err != nil { if err != nil {
log.Printf("[%s] 🤮 Cannot write the poll response: %s", m.Name, err) log.Printf("[%s] Could not write the map response: %s", m.Name, err)
} }
return true return true
case <-c.Request.Context().Done(): case <-c.Request.Context().Done():
log.Printf("[%s] 😥 The client has closed the connection", m.Name) log.Printf("[%s] The client has closed the connection", m.Name)
now := time.Now().UTC() now := time.Now().UTC()
m.LastSeen = &now m.LastSeen = &now
db.Save(&m) db.Save(&m)
h.pollMu.Lock() h.pollMu.Lock()
cancelKeepAlive <- []byte{} cancelKeepAlive <- []byte{}
delete(h.clientsPolling, m.ID) delete(h.clientsPolling, m.ID)
close(update)
h.pollMu.Unlock() h.pollMu.Unlock()
return false return false