From 9c9fde01602a3c4d0810e9aba634b0a9c4d2416d Mon Sep 17 00:00:00 2001
From: Kane York <rikingcoding@gmail.com>
Date: Wed, 8 Nov 2017 14:34:56 -0800
Subject: [PATCH] Use read deadlines instead of ping tickers

---
 socketserver/server/handlecore.go | 53 +++++++++++++++++++++++--------
 1 file changed, 39 insertions(+), 14 deletions(-)

diff --git a/socketserver/server/handlecore.go b/socketserver/server/handlecore.go
index 6e572819..197e6715 100644
--- a/socketserver/server/handlecore.go
+++ b/socketserver/server/handlecore.go
@@ -425,7 +425,24 @@ func runSocketReader(conn *websocket.Conn, client *ClientInfo, errorChan chan<-
 	defer close(errorChan)
 	defer close(clientChan)
 
-	for ; err == nil; messageType, packet, err = conn.ReadMessage() {
+	for {
+		conn.SetReadDeadline(time.Now().Add(1 * time.Minute))
+		messageType, packet, err = conn.ReadMessage()
+		// handle ReadDeadline by sending a ping
+		// writer loop handles repeated ping timeouts
+		if tmErr, ok := err.(interface {
+			Timeout() bool
+		}); ok && tmErr.Timeout() {
+			select {
+			case <-stoppedChan:
+				return
+			case clientChan <- ClientMessage{Command: "__readTimeout"}:
+			}
+			continue // re-set deadline and wait for pong packet
+		}
+		if err != nil {
+			break
+		}
 		if messageType == websocket.BinaryMessage {
 			err = &CloseGotBinaryMessage
 			break
@@ -448,23 +465,22 @@ func runSocketReader(conn *websocket.Conn, client *ClientInfo, errorChan chan<-
 		} else if msg.MessageID == 0 {
 			continue
 		}
+
 		select {
-		case clientChan <- msg:
 		case <-stoppedChan:
 			return
+		case clientChan <- msg:
 		}
 	}
 
 	select {
-	case errorChan <- err:
 	case <-stoppedChan:
+	case errorChan <- err:
 	}
 	// exit goroutine
 }
 
 func runSocketWriter(conn *websocket.Conn, client *ClientInfo, errorChan <-chan error, clientChan <-chan ClientMessage, serverMessageChan <-chan ClientMessage) websocket.CloseError {
-	pingTicker := time.NewTicker(1 * time.Minute)
-	defer pingTicker.Stop()
 	lastPacket := time.Now()
 
 	for {
@@ -485,12 +501,29 @@ func runSocketWriter(conn *websocket.Conn, client *ClientInfo, errorChan <-chan
 			}
 
 		case msg := <-clientChan:
+			if msg.Command == "__readTimeout" {
+				// generated on 60 seconds without a message
+				now := time.Now()
+				if lastPacket.Add(5 * time.Minute).Before(now) {
+					return CloseTimedOut
+				}
+				conn.WriteControl(
+					websocket.PingMessage,
+					[]byte(strconv.FormatInt(time.Now().Unix(), 10)),
+					getDeadline(),
+				)
+				continue
+			}
+
 			if !client.HelloOK && msg.Command != HelloCommand {
 				return CloseFirstMessageNotHello
 			}
 			lastPacket = time.Now()
+
 			if msg.Command == "__ping" {
-				continue // generated by server, not by client
+				// generated for PONG packets
+				// want this to run AFTER lastPacket was set
+				continue
 			}
 
 			for _, char := range msg.Command {
@@ -510,14 +543,6 @@ func runSocketWriter(conn *websocket.Conn, client *ClientInfo, errorChan <-chan
 			}
 			SendMessage(conn, msg)
 
-		case <-pingTicker.C:
-			now := time.Now()
-			if lastPacket.Add(5 * time.Minute).Before(now) {
-				return CloseTimedOut
-			} else if lastPacket.Add(1 * time.Minute).Before(now) {
-				conn.WriteControl(websocket.PingMessage, []byte(strconv.FormatInt(time.Now().Unix(), 10)), getDeadline())
-			}
-
 		case <-StopAcceptingConnectionsCh:
 			return CloseGoingAway
 		}