Changed: add possibility to set timeout, dial and transport

proxy.go

@@ -18,7 +18,7 @@ import (
 	"github.com/pkg/errors"
 )
 
-var errClientCertRequested = errors.New("tls: client cert authentication unsupported")
+var ErrClientCertRequested = errors.New("tls: client cert authentication unsupported")
 
 // defaultTimeout is the default value for reading from local connections.
 // By default we have no timeout.
@@ -28,6 +28,8 @@ const defaultTimeout = 0
 const dialTimeout = 30 * time.Second
 const tlsHandshakeTimeout = 10 * time.Second
 
+type DialFunc func(string, string) (net.Conn, error)
+
 // Proxy is a structure with the proxy server configuration and current state.
 type Proxy struct {
 	// addr is the address the proxy listens to.
@@ -42,7 +44,7 @@ type Proxy struct {
 
 	// dial is a function for creating net.Conn. Can be useful to override in
 	// unit-tests.
-	dial func(string, string) (net.Conn, error)
+	dial DialFunc
 
 	// timeout is the remote connection's read/write timeout.
 	timeout time.Duration
@@ -80,7 +82,7 @@ func NewProxy(config Config) *Proxy {
 					// http.Transport.RoundTrip method failed. In this case
 					// we'll receive the error and will be able to add the host
 					// to invalidTLSHosts.
-					return nil, errClientCertRequested
+					return nil, ErrClientCertRequested
 				},
 			},
 		},
@@ -102,6 +104,21 @@ func NewProxy(config Config) *Proxy {
 	return proxy
 }
 
+// SetTimeout set Timeout of the proxy (Deadline for processing the request).
+func (p *Proxy) SetTimeout(timeout time.Duration) {
+	p.timeout = timeout
+}
+
+// SetTransport set RoundTripper of the proxy.
+func (p *Proxy) SetTransport(transport *http.Transport) {
+	p.transport = transport
+}
+
+// SetDial set Dialer of the proxy.
+func (p *Proxy) SetDial(dial DialFunc) {
+	p.dial = dial
+}
+
 // Addr returns the address this proxy listens to.
 func (p *Proxy) Addr() (addr net.Addr) {
 	return p.addr
@@ -297,7 +314,7 @@ func (p *Proxy) handleRequest(ctx *Context) (err error) {
 			res = proxyutil.NewErrorResponse(session.req, err)
 
 			if strings.Contains(err.Error(), "x509: ") ||
-				strings.Contains(err.Error(), errClientCertRequested.Error()) {
+				errors.Is(err, ErrClientCertRequested) {
 				log.Printf("id=%s: adding %s to invalid TLS hosts due to: %v", session.ID(), session.req.Host, err)
 				p.invalidTLSHostsMu.Lock()
 				p.invalidTLSHosts[session.req.Host] = true
@@ -413,7 +430,7 @@ func (p *Proxy) handleTunnel(session *Session) (err error) {
 			// http.Transport.RoundTrip method failed. In this case we'll
 			// receive the error and will be able to add the host to
 			// invalidTLSHosts.
-			return nil, errClientCertRequested
+			return nil, ErrClientCertRequested
 		}
 
 		tlsConn := tls.Client(conn, &tls.Config{