Skip to content

Commit

Permalink
Update tlsClientHandshake
Browse files Browse the repository at this point in the history
  • Loading branch information
moredure committed Apr 7, 2022
1 parent 7a5afdd commit e985a73
Showing 1 changed file with 14 additions and 29 deletions.
43 changes: 14 additions & 29 deletions client.go
Expand Up @@ -1994,41 +1994,26 @@ func (c *HostClient) cachedTLSConfig(addr string) *tls.Config {
// ErrTLSHandshakeTimeout indicates there is a timeout from tls handshake.
var ErrTLSHandshakeTimeout = errors.New("tls handshake timed out")

var timeoutErrorChPool sync.Pool

func tlsClientHandshake(rawConn net.Conn, tlsConfig *tls.Config, timeout time.Duration) (net.Conn, error) {
tc := AcquireTimer(timeout)
defer ReleaseTimer(tc)

var ch chan error
chv := timeoutErrorChPool.Get()
if chv == nil {
chv = make(chan error)
}
ch = chv.(chan error)
defer timeoutErrorChPool.Put(chv)

func tlsClientHandshake(rawConn net.Conn, tlsConfig *tls.Config, deadline time.Time) (net.Conn, error) {
conn := tls.Client(rawConn, tlsConfig)

go func() {
ch <- conn.Handshake()
}()

select {
case <-tc.C:
_ = conn.SetReadDeadline(deadline)
_ = conn.SetWriteDeadline(deadline)
err := conn.Handshake()
if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
rawConn.Close()
<-ch
return nil, ErrTLSHandshakeTimeout
case err := <-ch:
if err != nil {
rawConn.Close()
return nil, err
}
return conn, nil
}
if err != nil {
rawConn.Close()
return nil, err
}
_ = conn.SetReadDeadline(time.Time{})
_ = conn.SetWriteDeadline(time.Time{})
return conn, nil
}

func dialAddr(addr string, dial DialFunc, dialDualStack, isTLS bool, tlsConfig *tls.Config, timeout time.Duration) (net.Conn, error) {
deadline := time.Now().Add(timeout)
if dial == nil {
if dialDualStack {
dial = DialDualStack
Expand All @@ -2049,7 +2034,7 @@ func dialAddr(addr string, dial DialFunc, dialDualStack, isTLS bool, tlsConfig *
if timeout == 0 {
return tls.Client(conn, tlsConfig), nil
}
return tlsClientHandshake(conn, tlsConfig, timeout)
return tlsClientHandshake(conn, tlsConfig, deadline)
}
return conn, nil
}
Expand Down

0 comments on commit e985a73

Please sign in to comment.