Prevent multiclient panic if underlying client is closed

Signed-off-by: Yilun <[email protected]>
nknorg · May 19, 2020 · af72bb3 · af72bb3
1 parent 5a50fba
commit af72bb3
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 13 deletions.
diff --git a/client.go b/client.go
@@ -458,17 +458,13 @@ func (c *Client) handleMessage(msgType int, data []byte) error {
 				data = nil
 			}
 
-			var msg *Message
-			switch payload.Type {
-			case payloads.BINARY, payloads.TEXT, payloads.SESSION:
-				msg = &Message{
-					Src:       inboundMsg.Src,
-					Data:      data,
-					Type:      int32(payload.Type),
-					Encrypted: payloadMsg.Encrypted,
-					MessageID: payload.MessageId,
-					NoReply:   payload.NoReply,
-				}
+			msg := &Message{
+				Src:       inboundMsg.Src,
+				Data:      data,
+				Type:      int32(payload.Type),
+				Encrypted: payloadMsg.Encrypted,
+				MessageID: payload.MessageId,
+				NoReply:   payload.NoReply,
 			}
 
 			if len(payload.ReplyToId) > 0 {

diff --git a/multiclient.go b/multiclient.go
@@ -155,12 +155,18 @@ func NewMultiClient(account *Account, baseIdentifier string, numSubClients int,
 
 			wg.Done()
 
-			node := <-client.OnConnect.C
+			node, ok := <-client.OnConnect.C
+			if !ok {
+				return
+			}
 			m.OnConnect.receive(node)
 
 			for {
 				select {
-				case msg := <-client.OnMessage.C:
+				case msg, ok := <-client.OnMessage.C:
+					if !ok {
+						return
+					}
 					if msg.Type == SessionType {
 						if !msg.Encrypted {
 							continue