Correctly cleanup/sever stream when forced to by remote peer

dominant-strategies · May 8, 2024 · aca1734 · aca1734
1 parent 76e58ad
commit aca1734
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/p2p/node/streamManager/streamManager.go b/p2p/node/streamManager/streamManager.go
@@ -81,8 +81,13 @@ func NewStreamManager(peerCount int, node quaiprotocol.QuaiP2PNode, host host.Ho
 	}, nil
 }
 
+// Expects a key as peerID and value of *streamWrapper
 func severStream(key interface{}, value interface{}) {
-	stream := value.(network.Stream)
+	wrappedStream, ok := value.(*streamWrapper)
+	if !ok {
+		return
+	}
+	stream := wrappedStream.stream
 	err := stream.Close()
 	if err != nil {
 		log.Global.WithField("err", err).Error("Failed to close stream")
@@ -98,10 +103,9 @@ func (sm *basicStreamManager) CloseStream(peerID p2p.PeerID) error {
 
 	wrappedStream, ok := sm.streamCache.Get(peerID)
 	if ok {
-		stream := wrappedStream.(*streamWrapper).stream
-		log.Global.WithField("peerID", peerID).Debug("Pruned connection with peer")
-		severStream(peerID, stream)
+		severStream(peerID, wrappedStream)
 		sm.streamCache.Remove(peerID)
+		log.Global.WithField("peerID", peerID).Debug("Pruned connection with peer")
 		return nil
 	}
 	return errStreamNotFound