From aa4fbf96175cf04da69bd1034c9943679298d5de Mon Sep 17 00:00:00 2001
From: Matt Caldwell <mcaldwell@endgame.com>
Date: Sun, 8 May 2016 17:54:50 -0400
Subject: [PATCH] Add hook for extending PongHandler

---
 session.go | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/session.go b/session.go
index d66d4b4..af51af8 100644
--- a/session.go
+++ b/session.go
@@ -74,14 +74,20 @@ 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()
+	})
+}
+
 func (s *Session) readPump() {
 	defer s.conn.Close()
 
 	s.conn.SetReadLimit(s.melody.Config.MaxMessageSize)
 	s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
 
-	s.conn.SetPongHandler(func(string) error {
-		s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
+	s.SetPongHandler(func() error {
 		return nil
 	})