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)