From 43f5a6e414c44f857a5b6de700db74d75d352c10 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9E=97=E5=BF=97=E5=AE=87?= <evanlinjin@pm.me>
Date: Sun, 2 Jun 2019 20:30:41 +0800
Subject: [PATCH] Fixed dmsg closing logic.

---
 pkg/dmsg/transport.go | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/pkg/dmsg/transport.go b/pkg/dmsg/transport.go
index 22fec920e1..86c0c9b848 100644
--- a/pkg/dmsg/transport.go
+++ b/pkg/dmsg/transport.go
@@ -56,8 +56,19 @@ func NewTransport(conn net.Conn, log *logging.Logger, local, remote cipher.PubKe
 
 func (c *Transport) close() (closed bool) {
 	c.doneOnce.Do(func() {
-		close(c.doneCh)
 		closed = true
+		close(c.doneCh)
+
+		// Kill all goroutines pushing to `c.readCh` before closing it.
+		// No more goroutines pushing to `c.readCh` should be created once `c.doneCh` is closed.
+		for {
+			select {
+			case <-c.readCh:
+			default:
+				close(c.readCh)
+				return
+			}
+		}
 	})
 	return closed
 }
@@ -122,7 +133,7 @@ func (c *Transport) InjectRead(f Frame) bool {
 	ok := c.injectRead(f)
 	if !ok {
 		c.close()
-		close(c.readCh)
+		//close(c.readCh)
 	}
 	return ok
 }