From 61206fdf42fcc284a01509494c4989016d2ded62 Mon Sep 17 00:00:00 2001 From: Mike Morris Date: Mon, 12 Aug 2019 12:47:02 -0400 Subject: [PATCH] snapshot: add TLS support to HalfCloser interface (#6216) Calls net.TCPConn.CloseWrite or mtls.Conn.CloseWrite, which was added in https://go-review.googlesource.com/c/go/+/31318/ --- agent/pool/pool.go | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/agent/pool/pool.go b/agent/pool/pool.go index e1ba1d8ec5..fd693514d1 100644 --- a/agent/pool/pool.go +++ b/agent/pool/pool.go @@ -2,6 +2,7 @@ package pool import ( "container/list" + "crypto/tls" "fmt" "io" "net" @@ -257,11 +258,8 @@ func (p *ConnPool) acquire(dc string, addr net.Addr, version int, useTLS bool) ( return nil, fmt.Errorf("rpc error: lead thread didn't get connection") } -// HalfCloser is an interface that exposes a TCP half-close. We need this -// because we want to expose the raw TCP connection underlying a TLS one in a -// way that's hard to screw up and use for anything else. There's a change -// brewing that will allow us to use the TLS connection for this instead - -// https://go-review.googlesource.com/#/c/25159/. +// HalfCloser is an interface that exposes a TCP half-close without exposing +// the underlying TLS or raw TCP connection. type HalfCloser interface { CloseWrite() error } @@ -296,11 +294,13 @@ func DialTimeoutWithRPCType(dc string, addr net.Addr, src *net.TCPAddr, timeout return nil, nil, err } - // Cast to TCPConn var hc HalfCloser + if tcp, ok := conn.(*net.TCPConn); ok { tcp.SetKeepAlive(true) tcp.SetNoDelay(true) + + // Expose TCPConn CloseWrite method on HalfCloser hc = tcp } @@ -319,6 +319,11 @@ func DialTimeoutWithRPCType(dc string, addr net.Addr, src *net.TCPAddr, timeout return nil, nil, err } conn = tlsConn + + // If this is a tls.Conn, expose HalfCloser to caller + if tlsConn, ok := conn.(*tls.Conn); ok { + hc = tlsConn + } } return conn, hc, nil