From 942b6083b0151d9756f990010e18540aa5925579 Mon Sep 17 00:00:00 2001
From: Steven Allen <steven@stebalien.com>
Date: Fri, 30 Jul 2021 12:03:30 -0700
Subject: [PATCH] fix: fix a map access race condition in the want index

---
 internal/decision/engine.go | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/internal/decision/engine.go b/internal/decision/engine.go
index c22a4d7f..31c50e3f 100644
--- a/internal/decision/engine.go
+++ b/internal/decision/engine.go
@@ -605,6 +605,7 @@ func (e *Engine) ReceiveFrom(from peer.ID, blks []blocks.Block) {
 
 	// Check each peer to see if it wants one of the blocks we received
 	var work bool
+	missingWants := make(map[peer.ID][]cid.Cid)
 	e.lock.RLock()
 	for _, b := range blks {
 		k := b.Cid()
@@ -613,7 +614,7 @@ func (e *Engine) ReceiveFrom(from peer.ID, blks []blocks.Block) {
 			ledger, ok := e.ledgerMap[p]
 			if !ok {
 				log.Errorw("failed to find peer in ledger", "peer", p)
-				e.peerLedger.CancelWant(p, k)
+				missingWants[p] = append(missingWants[p], k)
 				continue
 			}
 			ledger.lk.RLock()
@@ -621,7 +622,7 @@ func (e *Engine) ReceiveFrom(from peer.ID, blks []blocks.Block) {
 			ledger.lk.RUnlock()
 			if !ok { // should never happen
 				log.Errorw("wantlist index doesn't match peer's wantlist", "peer", p)
-				e.peerLedger.CancelWant(p, k)
+				missingWants[p] = append(missingWants[p], k)
 				continue
 			}
 			work = true
@@ -649,6 +650,30 @@ func (e *Engine) ReceiveFrom(from peer.ID, blks []blocks.Block) {
 	}
 	e.lock.RUnlock()
 
+	// If we found missing wants (e.g., because the peer disconnected, we have some races here)
+	// remove them from the list. Unfortunately, we still have to re-check because the user
+	// could have re-connected in the meantime.
+	if len(missingWants) > 0 {
+		e.lock.Lock()
+		for p, wl := range missingWants {
+			if ledger, ok := e.ledgerMap[p]; ok {
+				ledger.lk.RLock()
+				for _, k := range wl {
+					if _, has := ledger.WantListContains(k); has {
+						continue
+					}
+					e.peerLedger.CancelWant(p, k)
+				}
+				ledger.lk.RUnlock()
+			} else {
+				for _, k := range wl {
+					e.peerLedger.CancelWant(p, k)
+				}
+			}
+		}
+		e.lock.Unlock()
+	}
+
 	if work {
 		e.signalNewWork()
 	}