Fixing missing lock acquisition

session.go

@@ -386,6 +386,7 @@ func (s *Session) handleStreamMessage(hdr header) error {
 	if stream == nil {
 		// Drain any data on the wire
 		if hdr.MsgType() == typeData && hdr.Length() > 0 {
+			s.logger.Printf("[WARN] yamux: Discarding data for stream: %d", id)
 			if _, err := io.CopyN(ioutil.Discard, s.bufRead, int64(hdr.Length())); err != nil {
 				s.logger.Printf("[ERR] yamux: Failed to discard data: %v", err)
 				return nil
@@ -494,10 +495,8 @@ func (s *Session) incomingStream(id uint32) error {
 
 // closeStream is used to close a stream once both sides have
 // issued a close.
-func (s *Session) closeStream(id uint32, withLock bool) {
-	if !withLock {
-		s.streamLock.Lock()
-		defer s.streamLock.Unlock()
-	}
+func (s *Session) closeStream(id uint32) {
+	s.streamLock.Lock()
 	delete(s.streams, id)
+	s.streamLock.Unlock()
 }

stream.go

@@ -300,7 +300,7 @@ func (s *Stream) Close() error {
 	case streamLocalClose:
 	case streamRemoteClose:
 		s.state = streamClosed
-		s.session.closeStream(s.id, false)
+		s.session.closeStream(s.id)
 		goto SEND_CLOSE
 
 	case streamClosed:
@@ -346,7 +346,7 @@ func (s *Stream) processFlags(flags uint16) error {
 			s.notifyWaiting()
 		case streamLocalClose:
 			s.state = streamClosed
-			s.session.closeStream(s.id, true)
+			s.session.closeStream(s.id)
 			s.notifyWaiting()
 		default:
 			s.session.logger.Printf("[ERR] yamux: unexpected FIN flag in state %d", s.state)
@@ -354,7 +354,7 @@ func (s *Stream) processFlags(flags uint16) error {
 		}
 	} else if flags&flagRST == flagRST {
 		s.state = streamReset
-		s.session.closeStream(s.id, true)
+		s.session.closeStream(s.id)
 		s.notifyWaiting()
 	}
 	return nil