Add a initial read deadline to Socks5 Conn.

This allows us to actually implement the socks5 dial timeout without depend on the timeout setting on the socks5 proxy server.
This commit is contained in:
NI
2019-08-17 10:26:04 +08:00
parent fc9121e5c2
commit 541e4b9233

View File

@@ -24,6 +24,29 @@ import (
"golang.org/x/net/proxy" "golang.org/x/net/proxy"
) )
var (
emptyTime = time.Time{}
)
type socks5Conn struct {
net.Conn
initialReadDeadline time.Time
}
func (s *socks5Conn) Read(b []byte) (int, error) {
if s.initialReadDeadline != emptyTime {
s.SetReadDeadline(s.initialReadDeadline)
s.initialReadDeadline = emptyTime
defer s.SetReadDeadline(emptyTime)
}
rLen, rErr := s.Conn.Read(b)
return rLen, rErr
}
// BuildSocks5Dial builds a Socks5 dialer // BuildSocks5Dial builds a Socks5 dialer
func BuildSocks5Dial( func BuildSocks5Dial(
socks5Address string, userName string, password string) (Dial, error) { socks5Address string, userName string, password string) (Dial, error) {
@@ -41,15 +64,26 @@ func BuildSocks5Dial(
address string, address string,
timeout time.Duration, timeout time.Duration,
) (net.Conn, error) { ) (net.Conn, error) {
dial, dialErr := proxy.SOCKS5("tcp", socks5Address, auth, &net.Dialer{ dialCfg := net.Dialer{
Timeout: timeout, Timeout: timeout,
Deadline: time.Now().Add(timeout), Deadline: time.Now().Add(timeout),
}) }
dial, dialErr := proxy.SOCKS5("tcp", socks5Address, auth, &dialCfg)
if dialErr != nil { if dialErr != nil {
return nil, dialErr return nil, dialErr
} }
return dial.Dial(network, address) dialConn, dialErr := dial.Dial(network, address)
if dialErr != nil {
return nil, dialErr
}
return &socks5Conn{
Conn: dialConn,
initialReadDeadline: dialCfg.Deadline,
}, nil
}, nil }, nil
} }