From 846e1af9041ac4da0f43adafd83597a368fb4f68 Mon Sep 17 00:00:00 2001
From: Maxim Ivanov <hi+github@yamlcoder.me>
Date: Tue, 26 Jul 2022 14:58:41 +0100
Subject: [PATCH] Push net.Conn Deadline forward on each read.

SetDeadline sets absolute time after which all IO operations start to fail.

Previously deadline was pushed forward upon successful websocket message
receipt, which works for small messages, but fail to behave adequately
if large message is transmitted over a slow connection.

This change makes deadline to be pushed forward on every successful read,
allowing slow clients to complete transmission as long as they keep sending
bytes.
---
 server/session_ws.go | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/server/session_ws.go b/server/session_ws.go
index 3719b1dbe..3a2696c12 100644
--- a/server/session_ws.go
+++ b/server/session_ws.go
@@ -18,6 +18,8 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
+	"io/ioutil"
 	"net"
 	"sync"
 	"time"
@@ -170,6 +172,33 @@ func (s *sessionWS) Expiry() int64 {
 	return s.expiry
 }
 
+type ReaderFunc func(p []byte) (n int, err error)
+
+func (f ReaderFunc) Read(p []byte) (n int, err error) {
+	return f(p)
+}
+
+func (s *sessionWS) ReadMessageWithLivenessUpdate() (messageType int, p []byte, err error) {
+	var r io.Reader
+	messageType, r, err = s.conn.NextReader()
+	if err != nil {
+		return messageType, nil, err
+	}
+
+	// Wrap original Reader into ours, which resets ping timer on every successful
+	// read. This allows us to transmit large messages over slow connection and don't
+	// disconnect clients due to ping timeouts
+	LivenessUpdatingReader := func(p []byte) (n int, err error) {
+		n, err = r.Read(p)
+		if err != nil && n > 0 {
+			s.maybeResetPingTimer()
+		}
+		return
+	}
+	p, err = ioutil.ReadAll(ReaderFunc(LivenessUpdatingReader))
+	return
+}
+
 func (s *sessionWS) Consume() {
 	// Fire an event for session start.
 	if fn := s.runtime.EventSessionStart(); fn != nil {
@@ -195,7 +224,7 @@ func (s *sessionWS) Consume() {
 
 IncomingLoop:
 	for {
-		messageType, data, err := s.conn.ReadMessage()
+		messageType, data, err := s.ReadMessageWithLivenessUpdate()
 		if err != nil {
 			// Ignore "normal" WebSocket errors.
 			if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived) {
-- 
GitLab