PipePairWithTimeout

This commit is contained in:
Toby 2021-03-02 16:25:36 -08:00
parent 565d659338
commit 244d0d43a9
3 changed files with 53 additions and 100 deletions

View File

@ -3,7 +3,6 @@ package relay
import (
"github.com/tobyxdd/hysteria/pkg/core"
"github.com/tobyxdd/hysteria/pkg/utils"
"io"
"net"
"time"
)
@ -58,55 +57,8 @@ func (r *Relay) ListenAndServe() error {
return
}
defer rc.Close()
err = pipePair(c, rc, r.Timeout)
err = utils.PipePairWithTimeout(c, rc, r.Timeout)
r.ErrorFunc(c.RemoteAddr(), err)
}(c)
}
}
func pipePair(conn *net.TCPConn, stream io.ReadWriteCloser, timeout time.Duration) error {
errChan := make(chan error, 2)
// TCP to stream
go func() {
buf := make([]byte, utils.PipeBufferSize)
for {
if timeout != 0 {
_ = conn.SetDeadline(time.Now().Add(timeout))
}
rn, err := conn.Read(buf)
if rn > 0 {
_, err := stream.Write(buf[:rn])
if err != nil {
errChan <- err
return
}
}
if err != nil {
errChan <- err
return
}
}
}()
// Stream to TCP
go func() {
buf := make([]byte, utils.PipeBufferSize)
for {
rn, err := stream.Read(buf)
if rn > 0 {
_, err := conn.Write(buf[:rn])
if err != nil {
errChan <- err
return
}
if timeout != 0 {
_ = conn.SetDeadline(time.Now().Add(timeout))
}
}
if err != nil {
errChan <- err
return
}
}
}()
return <-errChan
}

View File

@ -7,7 +7,6 @@ import (
"github.com/tobyxdd/hysteria/pkg/acl"
"github.com/tobyxdd/hysteria/pkg/core"
"github.com/tobyxdd/hysteria/pkg/utils"
"io"
"strconv"
)
@ -173,7 +172,7 @@ func (s *Server) handleTCP(c *net.TCPConn, r *socks5.Request) error {
}
defer rc.Close()
_ = sendReply(c, socks5.RepSuccess)
closeErr = pipePair(c, rc, s.TCPTimeout)
closeErr = utils.PipePairWithTimeout(c, rc, s.TCPTimeout)
return nil
case acl.ActionProxy:
rc, err := s.HyClient.DialTCP(addr)
@ -184,7 +183,7 @@ func (s *Server) handleTCP(c *net.TCPConn, r *socks5.Request) error {
}
defer rc.Close()
_ = sendReply(c, socks5.RepSuccess)
closeErr = pipePair(c, rc, s.TCPTimeout)
closeErr = utils.PipePairWithTimeout(c, rc, s.TCPTimeout)
return nil
case acl.ActionBlock:
_ = sendReply(c, socks5.RepHostUnreachable)
@ -199,7 +198,7 @@ func (s *Server) handleTCP(c *net.TCPConn, r *socks5.Request) error {
}
defer rc.Close()
_ = sendReply(c, socks5.RepSuccess)
closeErr = pipePair(c, rc, s.TCPTimeout)
closeErr = utils.PipePairWithTimeout(c, rc, s.TCPTimeout)
return nil
default:
_ = sendReply(c, socks5.RepServerFailure)
@ -223,50 +222,3 @@ func parseRequestAddress(r *socks5.Request) (domain string, ip net.IP, port stri
return "", r.DstAddr, p, net.JoinHostPort(net.IP(r.DstAddr).String(), p)
}
}
func pipePair(conn *net.TCPConn, stream io.ReadWriteCloser, timeout time.Duration) error {
errChan := make(chan error, 2)
// TCP to stream
go func() {
buf := make([]byte, utils.PipeBufferSize)
for {
if timeout != 0 {
_ = conn.SetDeadline(time.Now().Add(timeout))
}
rn, err := conn.Read(buf)
if rn > 0 {
_, err := stream.Write(buf[:rn])
if err != nil {
errChan <- err
return
}
}
if err != nil {
errChan <- err
return
}
}
}()
// Stream to TCP
go func() {
buf := make([]byte, utils.PipeBufferSize)
for {
rn, err := stream.Read(buf)
if rn > 0 {
_, err := conn.Write(buf[:rn])
if err != nil {
errChan <- err
return
}
if timeout != 0 {
_ = conn.SetDeadline(time.Now().Add(timeout))
}
}
if err != nil {
errChan <- err
return
}
}
}()
return <-errChan
}

View File

@ -2,6 +2,8 @@ package utils
import (
"io"
"net"
"time"
)
const PipeBufferSize = 65536
@ -33,3 +35,50 @@ func Pipe2Way(rw1, rw2 io.ReadWriter) error {
// We only need the first error
return <-errChan
}
func PipePairWithTimeout(conn *net.TCPConn, stream io.ReadWriteCloser, timeout time.Duration) error {
errChan := make(chan error, 2)
// TCP to stream
go func() {
buf := make([]byte, PipeBufferSize)
for {
if timeout != 0 {
_ = conn.SetDeadline(time.Now().Add(timeout))
}
rn, err := conn.Read(buf)
if rn > 0 {
_, err := stream.Write(buf[:rn])
if err != nil {
errChan <- err
return
}
}
if err != nil {
errChan <- err
return
}
}
}()
// Stream to TCP
go func() {
buf := make([]byte, PipeBufferSize)
for {
rn, err := stream.Read(buf)
if rn > 0 {
_, err := conn.Write(buf[:rn])
if err != nil {
errChan <- err
return
}
if timeout != 0 {
_ = conn.SetDeadline(time.Now().Add(timeout))
}
}
if err != nil {
errChan <- err
return
}
}
}()
return <-errChan
}