From e21886f3a7b601a5509274fb14dcc66d859304a8 Mon Sep 17 00:00:00 2001
From: lesismal <40462947+lesismal@users.noreply.github.com>
Date: Thu, 20 May 2021 03:01:17 +0800
Subject: [PATCH] Fix concurrent panic

---
 melody.go  | 15 ++++++++-------
 session.go | 36 +++++++++++++++++++-----------------
 2 files changed, 27 insertions(+), 24 deletions(-)

diff --git a/melody.go b/melody.go
index d845876..6123e77 100644
--- a/melody.go
+++ b/melody.go
@@ -174,13 +174,14 @@ func (m *Melody) HandleRequestWithKeys(w http.ResponseWriter, r *http.Request, k
 	}
 
 	session := &Session{
-		Request: r,
-		Keys:    keys,
-		conn:    conn,
-		output:  make(chan *envelope, m.Config.MessageBufferSize),
-		melody:  m,
-		open:    true,
-		rwmutex: &sync.RWMutex{},
+		Request:    r,
+		Keys:       keys,
+		conn:       conn,
+		output:     make(chan *envelope, m.Config.MessageBufferSize),
+		outputDone: make(chan struct{}),
+		melody:     m,
+		open:       true,
+		rwmutex:    &sync.RWMutex{},
 	}
 
 	m.hub.register <- session
diff --git a/session.go b/session.go
index 3997cef..c34a785 100644
--- a/session.go
+++ b/session.go
@@ -11,13 +11,14 @@ import (
 
 // Session wrapper around websocket connections.
 type Session struct {
-	Request *http.Request
-	Keys    map[string]interface{}
-	conn    *websocket.Conn
-	output  chan *envelope
-	melody  *Melody
-	open    bool
-	rwmutex *sync.RWMutex
+	Request    *http.Request
+	Keys       map[string]interface{}
+	conn       *websocket.Conn
+	output     chan *envelope
+	outputDone chan struct{}
+	melody     *Melody
+	open       bool
+	rwmutex    *sync.RWMutex
 }
 
 func (s *Session) writeMessage(message *envelope) {
@@ -56,12 +57,13 @@ func (s *Session) closed() bool {
 }
 
 func (s *Session) close() {
-	if !s.closed() {
-		s.rwmutex.Lock()
-		s.open = false
+	s.rwmutex.Lock()
+	open := s.open
+	s.open = false
+	s.rwmutex.Unlock()
+	if open {
 		s.conn.Close()
-		close(s.output)
-		s.rwmutex.Unlock()
+		close(s.outputDone)
 	}
 }
 
@@ -76,11 +78,7 @@ func (s *Session) writePump() {
 loop:
 	for {
 		select {
-		case msg, ok := <-s.output:
-			if !ok {
-				break loop
-			}
-
+		case msg := <-s.output:
 			err := s.writeRaw(msg)
 
 			if err != nil {
@@ -101,6 +99,10 @@ loop:
 			}
 		case <-ticker.C:
 			s.ping()
+		case _, ok := <-s.outputDone:
+			if !ok {
+				break loop
+			}
 		}
 	}
 }