Merge pull request #1520 from tvoran/VAULT-1776/custom-dialer

Support custom transport dialer

config/transport.go

@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"runtime"
 	"time"
+
+	"github.com/hashicorp/consul-template/dependency"
 )
 
 const (
@@ -36,6 +38,11 @@ var (
 // TransportConfig is the configuration to tune low-level APIs for the
 // interactions on the wire.
 type TransportConfig struct {
+	// CustomDialer overrides the default net.Dial with a custom dialer. This is
+	// useful for instance with Vault Agent Templating to direct Consul Template
+	// requests through an internal cache.
+	CustomDialer dependency.TransportDialer `mapstructure:"-"`
+
 	// DialKeepAlive is the amount of time for keep-alives.
 	DialKeepAlive *time.Duration `mapstructure:"dial_keep_alive"`
 
@@ -75,6 +82,7 @@ func (c *TransportConfig) Copy() *TransportConfig {
 
 	var o TransportConfig
 
+	o.CustomDialer = c.CustomDialer
 	o.DialKeepAlive = c.DialKeepAlive
 	o.DialTimeout = c.DialTimeout
 	o.DisableKeepAlives = c.DisableKeepAlives
@@ -104,6 +112,10 @@ func (c *TransportConfig) Merge(o *TransportConfig) *TransportConfig {
 
 	r := c.Copy()
 
+	if o.CustomDialer != nil {
+		r.CustomDialer = o.CustomDialer
+	}
+
 	if o.DialKeepAlive != nil {
 		r.DialKeepAlive = o.DialKeepAlive
 	}

config/transport_test.go

@@ -2,6 +2,7 @@ package config
 
 import (
 	"fmt"
+	"net"
 	"reflect"
 	"testing"
 	"time"
@@ -33,6 +34,17 @@ func TestTransportConfig_Copy(t *testing.T) {
 				TLSHandshakeTimeout: TimeDuration(30 * time.Second),
 			},
 		},
+		{
+			"same_enabled_custom_dialer",
+			&TransportConfig{
+				CustomDialer:        &net.Dialer{Timeout: 10 * time.Second},
+				DisableKeepAlives:   Bool(true),
+				IdleConnTimeout:     TimeDuration(40 * time.Second),
+				MaxIdleConns:        Int(150),
+				MaxIdleConnsPerHost: Int(15),
+				TLSHandshakeTimeout: TimeDuration(30 * time.Second),
+			},
+		},
 	}
 
 	for i, tc := range cases {
@@ -245,6 +257,30 @@ func TestTransportConfig_Merge(t *testing.T) {
 			&TransportConfig{TLSHandshakeTimeout: TimeDuration(10 * time.Second)},
 			&TransportConfig{TLSHandshakeTimeout: TimeDuration(10 * time.Second)},
 		},
+		{
+			"custom_transport_dialer",
+			&TransportConfig{CustomDialer: &net.Dialer{Timeout: 10 * time.Second}},
+			&TransportConfig{CustomDialer: &net.Dialer{Timeout: 20 * time.Second}},
+			&TransportConfig{CustomDialer: &net.Dialer{Timeout: 20 * time.Second}},
+		},
+		{
+			"custom_transport_dialer_empty_one",
+			&TransportConfig{CustomDialer: &net.Dialer{Timeout: 10 * time.Second}},
+			&TransportConfig{},
+			&TransportConfig{CustomDialer: &net.Dialer{Timeout: 10 * time.Second}},
+		},
+		{
+			"custom_transport_dialer_empty_two",
+			&TransportConfig{},
+			&TransportConfig{CustomDialer: &net.Dialer{Timeout: 10 * time.Second}},
+			&TransportConfig{CustomDialer: &net.Dialer{Timeout: 10 * time.Second}},
+		},
+		{
+			"custom_transport_dialer_same",
+			&TransportConfig{CustomDialer: &net.Dialer{Timeout: 10 * time.Second}},
+			&TransportConfig{CustomDialer: &net.Dialer{Timeout: 10 * time.Second}},
+			&TransportConfig{CustomDialer: &net.Dialer{Timeout: 10 * time.Second}},
+		},
 	}
 
 	for i, tc := range cases {

dependency/client_set.go

@@ -1,6 +1,7 @@
 package dependency
 
 import (
+	"context"
 	"crypto/tls"
 	"fmt"
 	"log"
@@ -35,6 +36,16 @@ type vaultClient struct {
 	httpClient *http.Client
 }
 
+// TransportDialer is an interface that allows passing a custom dialer function
+// to an HTTP client's transport config
+type TransportDialer interface {
+	// Dial is intended to match https://pkg.go.dev/net#Dialer.Dial
+	Dial(network, address string) (net.Conn, error)
+
+	// DialContext is intended to match https://pkg.go.dev/net#Dialer.DialContext
+	DialContext(ctx context.Context, network, address string) (net.Conn, error)
+}
+
 // CreateConsulClientInput is used as input to the CreateConsulClient function.
 type CreateConsulClientInput struct {
 	Address      string
@@ -74,6 +85,7 @@ type CreateVaultClientInput struct {
 	SSLCAPath   string
 	ServerName  string
 
+	TransportCustomDialer        TransportDialer
 	TransportDialKeepAlive       time.Duration
 	TransportDialTimeout         time.Duration
 	TransportDisableKeepAlives   bool
@@ -202,12 +214,19 @@ func (c *ClientSet) CreateVaultClient(i *CreateVaultClientInput) error {
 	}
 
 	// This transport will attempt to keep connections open to the Vault server.
+	var dialer TransportDialer
+	dialer = &net.Dialer{
+		Timeout:   i.TransportDialTimeout,
+		KeepAlive: i.TransportDialKeepAlive,
+	}
+
+	if i.TransportCustomDialer != nil {
+		dialer = i.TransportCustomDialer
+	}
+
 	transport := &http.Transport{
-		Proxy: http.ProxyFromEnvironment,
-		Dial: (&net.Dialer{
-			Timeout:   i.TransportDialTimeout,
-			KeepAlive: i.TransportDialKeepAlive,
-		}).Dial,
+		Proxy:               http.ProxyFromEnvironment,
+		Dial:                dialer.Dial,
 		DisableKeepAlives:   i.TransportDisableKeepAlives,
 		MaxIdleConns:        i.TransportMaxIdleConns,
 		IdleConnTimeout:     i.TransportIdleConnTimeout,

manager/runner.go

@@ -1284,6 +1284,7 @@ func newClientSet(c *config.Config) (*dep.ClientSet, error) {
 		SSLCACert:                    config.StringVal(c.Vault.SSL.CaCert),
 		SSLCAPath:                    config.StringVal(c.Vault.SSL.CaPath),
 		ServerName:                   config.StringVal(c.Vault.SSL.ServerName),
+		TransportCustomDialer:        c.Vault.Transport.CustomDialer,
 		TransportDialKeepAlive:       config.TimeDurationVal(c.Vault.Transport.DialKeepAlive),
 		TransportDialTimeout:         config.TimeDurationVal(c.Vault.Transport.DialTimeout),
 		TransportDisableKeepAlives:   config.BoolVal(c.Vault.Transport.DisableKeepAlives),