From ba467e8a324b56557f1c1e26317510df729a8faa Mon Sep 17 00:00:00 2001
From: Andy Wang <cbeuw.andy@gmail.com>
Date: Mon, 19 Aug 2019 11:31:56 +0100
Subject: [PATCH] Only store stream in map after Accept

---
 internal/multiplex/session.go | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/internal/multiplex/session.go b/internal/multiplex/session.go
index e050108..ddede6a 100644
--- a/internal/multiplex/session.go
+++ b/internal/multiplex/session.go
@@ -123,6 +123,9 @@ func (sesh *Session) Accept() (net.Conn, error) {
 	if stream == nil {
 		return nil, ErrBrokenSession
 	}
+	sesh.streamsM.Lock()
+	sesh.streams[stream.id] = stream
+	sesh.streamsM.Unlock()
 	log.Tracef("stream %v of session %v accepted", stream.id, sesh.id)
 	return stream, nil
 }
@@ -161,7 +164,6 @@ func (sesh *Session) recvDataFromRemote(data []byte) error {
 			connId, _ := sesh.sb.assignRandomConn()
 			// we ignore the error here. If the switchboard is broken, it will be reflected upon stream.Write
 			stream = makeStream(sesh, frame.StreamID, connId)
-			sesh.streams[frame.StreamID] = stream
 			sesh.acceptCh <- stream
 			stream.writeFrame(frame)
 			return nil