From 6274ab9053ca36ed554d3c00fcc8b0d4de5624a0 Mon Sep 17 00:00:00 2001
From: Bruce Marriner <bruce@sqls.net>
Date: Mon, 11 Jan 2016 14:05:22 -0600
Subject: [PATCH] Hackish fix so we don't try to close nil or closed channels.

---
 wsapi.go | 42 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 40 insertions(+), 2 deletions(-)

diff --git a/wsapi.go b/wsapi.go
index 33d37d8..fc50953 100644
--- a/wsapi.go
+++ b/wsapi.go
@@ -114,7 +114,33 @@ func (s *Session) Listen() (err error) {
 	s.listenChan = make(chan struct{})
 	s.listenLock.Unlock()
 
-	defer close(s.heartbeatChan)
+	// this is ugly.
+	defer func() {
+		if s.listenChan == nil {
+			return
+		}
+		select {
+		case <-s.listenChan:
+			break
+		default:
+			close(s.listenChan)
+		}
+		s.listenChan = nil
+	}()
+
+	// this is ugly.
+	defer func() {
+		if s.heartbeatChan == nil {
+			return
+		}
+		select {
+		case <-s.heartbeatChan:
+			break
+		default:
+			close(s.heartbeatChan)
+		}
+		s.listenChan = nil
+	}()
 
 	for {
 		messageType, message, err := s.wsConn.ReadMessage()
@@ -506,7 +532,19 @@ func (s *Session) Heartbeat(i time.Duration) {
 	s.heartbeatChan = make(chan struct{})
 	s.heartbeatLock.Unlock()
 
-	defer close(s.heartbeatChan)
+	// this is ugly.
+	defer func() {
+		if s.heartbeatChan == nil {
+			return
+		}
+		select {
+		case <-s.heartbeatChan:
+			break
+		default:
+			close(s.heartbeatChan)
+		}
+		s.listenChan = nil
+	}()
 
 	// send first heartbeat immediately because lag could put the
 	// first heartbeat outside the required heartbeat interval window