Merge branch 'fix-termination-of-long-downloads' into notsure2

This commit is contained in:
notsure2 2020-12-09 05:09:40 +02:00
commit c87405377c
1 changed files with 5 additions and 3 deletions

View File

@ -85,14 +85,17 @@ func RouteTCP(listener net.Listener, streamTimeout time.Duration, newSeshFunc fu
if sesh == nil || sesh.IsClosed() || sesh.Singleplex { if sesh == nil || sesh.IsClosed() || sesh.Singleplex {
sesh = newSeshFunc() 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) data := make([]byte, 10240)
_ = localConn.SetReadDeadline(time.Now().Add(streamTimeout))
i, err := io.ReadAtLeast(localConn, data, 1) i, err := io.ReadAtLeast(localConn, data, 1)
if err != nil { if err != nil {
log.Errorf("Failed to read first packet from proxy client: %v", err) log.Errorf("Failed to read first packet from proxy client: %v", err)
localConn.Close() localConn.Close()
return return
} }
var zeroTime time.Time
_ = localConn.SetReadDeadline(zeroTime)
stream, err := sesh.OpenStream() stream, err := sesh.OpenStream()
if err != nil { if err != nil {
@ -112,7 +115,6 @@ func RouteTCP(listener net.Listener, streamTimeout time.Duration, newSeshFunc fu
return return
} }
stream.SetReadFromTimeout(streamTimeout) // if localConn hasn't sent anything to stream to a period of time, stream closes
go func() { go func() {
if _, err := common.Copy(localConn, stream); err != nil { if _, err := common.Copy(localConn, stream); err != nil {
log.Tracef("copying stream to proxy client: %v", err) 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 { if _, err = common.Copy(stream, localConn); err != nil {
log.Tracef("copying proxy client to stream: %v", err) log.Tracef("copying proxy client to stream: %v", err)
} }
}(sesh, localConn) }(sesh, localConn, streamTimeout)
} }
} }