From 3fb44fab6c050fda1df3a374c33053dcb5f578ba Mon Sep 17 00:00:00 2001
From: Matt Caldwell <mcaldwell@endgame.com>
Date: Sun, 8 May 2016 20:21:27 -0400
Subject: [PATCH] Use a mutex to ensure SetPongHandler atomicity

---
 session.go | 25 ++++++++++++++++++-------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/session.go b/session.go
index af51af8..24fb950 100644
--- a/session.go
+++ b/session.go
@@ -4,6 +4,7 @@ import (
 	"errors"
 	"github.com/gorilla/websocket"
 	"net/http"
+	"sync"
 	"time"
 )
 
@@ -11,6 +12,7 @@ import (
 type Session struct {
 	Request *http.Request
 	conn    *websocket.Conn
+	mutex   sync.Mutex
 	output  chan *envelope
 	melody  *Melody
 }
@@ -75,10 +77,14 @@ loop:
 }
 
 func (s *Session) SetPongHandler(f func() error) {
-	s.conn.SetPongHandler(func(string) error {
-		s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
-		return f()
-	})
+	s.mutex.Lock()
+	{
+		s.conn.SetPongHandler(func(string) error {
+			s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
+			return f()
+		})
+	}
+	s.mutex.Unlock()
 }
 
 func (s *Session) readPump() {
@@ -87,9 +93,14 @@ func (s *Session) readPump() {
 	s.conn.SetReadLimit(s.melody.Config.MaxMessageSize)
 	s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
 
-	s.SetPongHandler(func() error {
-		return nil
-	})
+	s.mutex.Lock()
+	{
+		s.conn.SetPongHandler(func(string) error {
+			s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
+			return nil
+		})
+	}
+	s.mutex.Unlock()
 
 	for {
 		t, message, err := s.conn.ReadMessage()