From e77fd4c446f9e0933008c28032e594eaf55631d8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=8D=E7=A1=AE=E5=AE=9A?=
 <35424927+notsure2@users.noreply.github.com>
Date: Sat, 12 Dec 2020 19:00:46 +0200
Subject: [PATCH] Fix regression: termination of long downloads after
 StreamTimeout seconds (#141)

* Fix termination of long downloads after StreamTimeout seconds.

- Even if not broadcasting in a loop, we still need to update the read deadline.
- Don't enforce the timeout after the first data is written.

* When timeout no longer needs to be enforced, no need to schedule a broadcast.

* Fix Cloak client. Don't enforce read deadline after first read.

* Enforce StreamTimeout on the initial bytes sent by localConn only.

* Revert changes to multiplex module. Remove timeout from caller.
---
 internal/client/piper.go      | 8 +++++---
 internal/server/dispatcher.go | 2 --
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/internal/client/piper.go b/internal/client/piper.go
index fd1746a..295e3e8 100644
--- a/internal/client/piper.go
+++ b/internal/client/piper.go
@@ -85,14 +85,17 @@ func RouteTCP(listener net.Listener, streamTimeout time.Duration, newSeshFunc fu
 		if sesh == nil || sesh.IsClosed() || sesh.Singleplex {
 			sesh = newSeshFunc()
 		}
-		go func(sesh *mux.Session, localConn net.Conn) {
+		go func(sesh *mux.Session, localConn net.Conn, timeout time.Duration) {
 			data := make([]byte, 10240)
+			_ = localConn.SetReadDeadline(time.Now().Add(streamTimeout))
 			i, err := io.ReadAtLeast(localConn, data, 1)
 			if err != nil {
 				log.Errorf("Failed to read first packet from proxy client: %v", err)
 				localConn.Close()
 				return
 			}
+			var zeroTime time.Time
+			_ = localConn.SetReadDeadline(zeroTime)
 
 			stream, err := sesh.OpenStream()
 			if err != nil {
@@ -112,7 +115,6 @@ func RouteTCP(listener net.Listener, streamTimeout time.Duration, newSeshFunc fu
 				return
 			}
 
-			stream.SetReadFromTimeout(streamTimeout) // if localConn hasn't sent anything to stream to a period of time, stream closes
 			go func() {
 				if _, err := common.Copy(localConn, stream); err != nil {
 					log.Tracef("copying stream to proxy client: %v", err)
@@ -121,7 +123,7 @@ func RouteTCP(listener net.Listener, streamTimeout time.Duration, newSeshFunc fu
 			if _, err = common.Copy(stream, localConn); err != nil {
 				log.Tracef("copying proxy client to stream: %v", err)
 			}
-		}(sesh, localConn)
+		}(sesh, localConn, streamTimeout)
 	}
 
 }
diff --git a/internal/server/dispatcher.go b/internal/server/dispatcher.go
index 56a556f..4fa0698 100644
--- a/internal/server/dispatcher.go
+++ b/internal/server/dispatcher.go
@@ -275,8 +275,6 @@ func serveSession(sesh *mux.Session, ci ClientInfo, user *ActiveUser, sta *State
 		}
 		log.Tracef("%v endpoint has been successfully connected", ci.ProxyMethod)
 
-		// if stream has nothing to send to proxy server for sta.Timeout period of time, stream will return error
-		newStream.(*mux.Stream).SetWriteToTimeout(sta.Timeout)
 		go func() {
 			if _, err := common.Copy(localConn, newStream); err != nil {
 				log.Tracef("copying stream to proxy server: %v", err)