From fa8dbbf18e05f01b249e900e2f7c2ad8e746f2fd Mon Sep 17 00:00:00 2001
From: Kale Blankenship <kale@lemnisys.com>
Date: Wed, 7 Feb 2018 21:26:49 -0800
Subject: [PATCH] Fix conn.mux session lookup.

conn.mux should always lookup session by RemoteChannel when a Begin is received
instead of only when lookup by Channel fails.
---
 conn.go | 29 +++++++++++++++++------------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/conn.go b/conn.go
index 34f3414..f75787a 100644
--- a/conn.go
+++ b/conn.go
@@ -314,24 +314,29 @@ func (c *conn) mux() {
 
 		// new frame from connReader
 		case fr := <-c.rxFrame:
-			// lookup session and send to Session.mux
-			session, ok := sessionsByRemoteChannel[fr.channel]
-			if !ok {
-				// if this is a begin, RemoteChannel should be used
-				begin, ok := fr.body.(*performBegin)
-				if !ok {
-					c.err = errorErrorf("unexpected frame: %#v", fr.body)
-					continue
-				}
+			var (
+				session *Session
+				ok      bool
+			)
 
-				session, ok = sessionsByChannel[begin.RemoteChannel]
+			switch body := fr.body.(type) {
+			// RemoteChannel should be used when frame is Begin
+			case *performBegin:
+				session, ok = sessionsByChannel[body.RemoteChannel]
 				if !ok {
-					c.err = errorErrorf("unexpected frame: %#v", fr.body)
-					continue
+					break
 				}
 
 				session.remoteChannel = fr.channel
 				sessionsByRemoteChannel[fr.channel] = session
+
+			default:
+				session, ok = sessionsByRemoteChannel[fr.channel]
+			}
+
+			if !ok {
+				c.err = errorErrorf("unexpected frame: %#v", fr.body)
+				continue
 			}
 
 			select {