From 44f212e3ea7403128703ace6b88cd52e97e2d1e3 Mon Sep 17 00:00:00 2001
From: Denis Arh <denis.arh@kendu.si>
Date: Tue, 23 Oct 2018 13:19:32 +0200
Subject: [PATCH] Fix websocket disconnection due to misconfiguration

---
 internal/config/websocket.go |  2 +-
 sam/websocket/session.go     | 28 +++++++++++++++++++---------
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/internal/config/websocket.go b/internal/config/websocket.go
index c3015e4d7..a27cd46eb 100644
--- a/internal/config/websocket.go
+++ b/internal/config/websocket.go
@@ -25,6 +25,6 @@ func (*Websocket) Init(prefix ...string) *Websocket {
 	websocket = new(Websocket)
 	websocket.Timeout = 15 * time.Second
 	websocket.PingTimeout = 120 * time.Second
-	websocket.PingPeriod = (websocket.PingTimeout * 10) / 9
+	websocket.PingPeriod = (websocket.PingTimeout * 9) / 10
 	return websocket
 }
diff --git a/sam/websocket/session.go b/sam/websocket/session.go
index 7aa92f9eb..b0027d9af 100644
--- a/sam/websocket/session.go
+++ b/sam/websocket/session.go
@@ -111,17 +111,20 @@ func (sess *Session) Close() {
 	store.Delete(sess.id)
 }
 
-func (sess *Session) readLoop() error {
+func (sess *Session) readLoop() (err error) {
 	defer func() {
 		log.Println("serveWebsocket - stop")
 		sess.Close()
 	}()
 
-	sess.conn.SetReadDeadline(time.Now().Add(sess.config.Websocket.PingTimeout))
+	if err = sess.conn.SetReadDeadline(time.Now().Add(sess.config.Websocket.PingTimeout)); err != nil {
+		return
+	}
+
 	sess.conn.SetPongHandler(func(string) error {
-		sess.conn.SetReadDeadline(time.Now().Add(sess.config.Websocket.PingTimeout))
-		return nil
+		return sess.conn.SetReadDeadline(time.Now().Add(sess.config.Websocket.PingTimeout))
 	})
+
 	sess.remoteAddr = sess.conn.RemoteAddr().String()
 
 	for {
@@ -145,16 +148,23 @@ func (sess *Session) writeLoop() error {
 		sess.Close() // break readLoop
 	}()
 
-	write := func(msg []byte) error {
-		sess.conn.SetWriteDeadline(time.Now().Add(sess.config.Websocket.Timeout))
+	write := func(msg []byte) (err error) {
+		if err = sess.conn.SetWriteDeadline(time.Now().Add(sess.config.Websocket.Timeout)); err != nil {
+			return
+		}
+
 		if msg != nil {
 			return sess.conn.WriteMessage(websocket.TextMessage, msg)
 		}
-		return nil
+
+		return
 	}
 
-	ping := func() error {
-		sess.conn.SetWriteDeadline(time.Now().Add(sess.config.Websocket.Timeout))
+	ping := func() (err error) {
+		if err = sess.conn.SetWriteDeadline(time.Now().Add(sess.config.Websocket.Timeout)); err != nil {
+			return
+		}
+
 		return sess.conn.WriteMessage(websocket.PingMessage, nil)
 	}