Prevent deadlock with closeStream race

stream.go

@@ -267,6 +267,7 @@ func (s *Stream) sendClose() error {
 
 // Close is used to close the stream
 func (s *Stream) Close() error {
+	closeStream := false
 	s.stateLock.Lock()
 	switch s.state {
 	// Opened means we need to signal a close
@@ -281,7 +282,7 @@ func (s *Stream) Close() error {
 	case streamLocalClose:
 	case streamRemoteClose:
 		s.state = streamClosed
-		s.session.closeStream(s.id)
+		closeStream = true
 		goto SEND_CLOSE
 
 	case streamClosed:
@@ -295,6 +296,9 @@ SEND_CLOSE:
 	s.stateLock.Unlock()
 	s.sendClose()
 	s.notifyWaiting()
+	if closeStream {
+		s.session.closeStream(s.id)
+	}
 	return nil
 }
 
@@ -309,6 +313,14 @@ func (s *Stream) forceClose() {
 // processFlags is used to update the state of the stream
 // based on set flags, if any. Lock must be held
 func (s *Stream) processFlags(flags uint16) error {
+	// Close the stream without holding the state lock
+	closeStream := false
+	defer func() {
+		if closeStream {
+			s.session.closeStream(s.id)
+		}
+	}()
+
 	s.stateLock.Lock()
 	defer s.stateLock.Unlock()
 	if flags&flagACK == flagACK {
@@ -327,7 +339,7 @@ func (s *Stream) processFlags(flags uint16) error {
 			s.notifyWaiting()
 		case streamLocalClose:
 			s.state = streamClosed
-			s.session.closeStream(s.id)
+			closeStream = true
 			s.notifyWaiting()
 		default:
 			s.session.logger.Printf("[ERR] yamux: unexpected FIN flag in state %d", s.state)
@@ -335,7 +347,7 @@ func (s *Stream) processFlags(flags uint16) error {
 		}
 	} else if flags&flagRST == flagRST {
 		s.state = streamReset
-		s.session.closeStream(s.id)
+		closeStream = true
 		s.notifyWaiting()
 	}
 	return nil