From ae889c5259380a59227aaa2f11ce45bab46ac78b Mon Sep 17 00:00:00 2001
From: Abhilash Gnan <abhilashgnan@gmail.com>
Date: Thu, 21 May 2020 09:16:06 +0200
Subject: [PATCH] fix blocking select on donec

Signed-off-by: Abhilash Gnan <abhilashgnan@gmail.com>
---
 cmux.go | 42 +++++++++++++++++++++++++++---------------
 1 file changed, 27 insertions(+), 15 deletions(-)

diff --git a/cmux.go b/cmux.go
index 0d4d624..adf268a 100644
--- a/cmux.go
+++ b/cmux.go
@@ -15,6 +15,7 @@
 package cmux
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"net"
@@ -61,6 +62,9 @@ func (e errListenerClosed) Timeout() bool   { return false }
 // listener is closed.
 var ErrListenerClosed = errListenerClosed("mux: listener closed")
 
+// ErrServerClosed is returned from muxListener.Accept when mux server is closed.
+var ErrServerClosed = errors.New("mux: server closed")
+
 // for readability of readTimeout
 var noTimeout time.Duration
 
@@ -148,7 +152,7 @@ func (m *cMux) Serve() error {
 	var wg sync.WaitGroup
 
 	defer func() {
-		m.closeDoneChan()
+		m.closeDoneChans()
 		wg.Wait()
 
 		for _, sl := range m.sls {
@@ -161,14 +165,6 @@ func (m *cMux) Serve() error {
 	}()
 
 	for {
-		select {
-		case <-m.donec:
-			// cmux was closed with cmux.Close()
-			return nil
-		default:
-			// do nothing
-		}
-
 		c, err := m.root.Accept()
 		if err != nil {
 			if !m.handleErr(err) {
@@ -215,16 +211,24 @@ func (m *cMux) serve(c net.Conn, donec <-chan struct{}, wg *sync.WaitGroup) {
 }
 
 func (m *cMux) Close() {
-	m.closeDoneChan()
+	m.closeDoneChans()
 }
 
-func (m *cMux) closeDoneChan() {
+func (m *cMux) closeDoneChans() {
 	select {
 	case <-m.donec:
 		// Already closed. Don't close again
 	default:
 		close(m.donec)
 	}
+	for _, sl := range m.sls {
+		select {
+		case <-sl.l.donec:
+			// Already closed. Don't close again
+		default:
+			close(sl.l.donec)
+		}
+	}
 }
 
 func (m *cMux) HandleError(h ErrorHandler) {
@@ -246,14 +250,22 @@ func (m *cMux) handleErr(err error) bool {
 type muxListener struct {
 	net.Listener
 	connc chan net.Conn
+	donec chan struct{}
 }
 
 func (l muxListener) Accept() (net.Conn, error) {
-	c, ok := <-l.connc
-	if !ok {
-		return nil, ErrListenerClosed
+	select {
+	case c, ok := <-l.connc:
+		if !ok {
+			return nil, ErrListenerClosed
+		}
+		return c, nil
+	case <-l.donec:
+		return nil, ErrServerClosed
+	default:
+		// do nothing
 	}
-	return c, nil
+	return nil, ErrServerClosed
 }
 
 // MuxConn wraps a net.Conn and provides transparent sniffing of connection data.