Allow for local dns resolution with a custom dialer

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## 1.2.0
+
+### Features
+
+* A connector's dialer can now be used to resolve DNS if the dialer implements the `HostDialer` interface
+
 ## 1.0.0
 
 ### Features

README.md

@@ -72,6 +72,10 @@ Other supported formats are listed below.
 
 If no pipe name can be derived from the DSN, connection attempts will first query the SQL Browser service to find the pipe name for the instance.
 
+### DNS Resolution through a Custom Dialer
+
+Custom Dialers can be used to resolve DNS if the Connection's Dialer implements the `HostDialer` interface. This is helpful when the dialer is proxying requests to a different, private network and the DNS record is local to the private network.
+
 ### Protocol configuration
 
 To force a specific protocol for the connection there two several options:

mssql.go

@@ -194,7 +194,9 @@ type Connector struct {
 	// SessionInitSQL is empty.
 	SessionInitSQL string
 
-	// Dialer sets a custom dialer for all network operations.
+	// Dialer sets a custom dialer for all network operations, except DNS resolution unless
+	// the dialer implements the HostDialer.
+	//
 	// If Dialer is not set, normal net dialers are used.
 	Dialer Dialer
 }
@@ -203,6 +205,13 @@ type Dialer interface {
 	DialContext(ctx context.Context, network string, addr string) (net.Conn, error)
 }
 
+// HostDialer should be used if the dialer is proxying requests to a different network
+// and DNS should be resolved in that other network
+type HostDialer interface {
+	Dialer
+	HostName() string
+}
+
 func (c *Connector) getDialer(p *msdsn.Config) Dialer {
 	if c != nil && c.Dialer != nil {
 		return c.Dialer

protocol.go

@@ -51,6 +51,14 @@ func (t tcpDialer) DialSqlConnection(ctx context.Context, c *Connector, p *msdsn
 	var ips []net.IP
 	ip := net.ParseIP(p.Host)
 	if ip == nil {
+		// if the custom dialer is a host dialer, the DNS is resolved within the network
+		// the dialer is sending the request to, rather than the one the driver is running on
+		d := c.getDialer(p)
+		if _, ok := d.(HostDialer); ok {
+			addr := net.JoinHostPort(p.Host, strconv.Itoa(int(resolveServerPort(p.Port))))
+			return d.DialContext(ctx, "tcp", addr)
+		}
+
 		ips, err = net.LookupIP(p.Host)
 		if err != nil {
 			return

tds_login_test.go

@@ -43,6 +43,19 @@ func (d *MockTransportDialer) DialContext(ctx context.Context, network string, a
 	return d.client, nil
 }
 
+type MockHostTransportDialer struct {
+	Dialer *MockTransportDialer
+	Host   string
+}
+
+func (m MockHostTransportDialer) DialContext(ctx context.Context, network string, addr string) (conn net.Conn, err error) {
+	return m.Dialer.DialContext(ctx, network, addr)
+}
+
+func (m MockHostTransportDialer) HostName() string {
+	return m.Host
+}
+
 func testLoginSequenceServer(result chan error, conn net.Conn, expectedPackets, responsePackets []string) {
 	defer func() {
 		conn.Close()

tds_test.go

@@ -971,3 +971,56 @@ func versionToHexString(v uint32) string {
 	binary.LittleEndian.PutUint32(b, v)
 	return hex.EncodeToString(b)
 }
+
+func TestDialSqlConnectionCustomDialer(t *testing.T) {
+	tl := testLogger{t: t}
+	defer tl.StopLogging()
+	SetLogger(&tl)
+
+	params := msdsn.Config{
+		Host: "nonexistant-dns.svc.cluster.local",
+	}
+	connector, err := NewConnector(params.URL().String())
+	if err != nil {
+		t.Error(err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
+	defer cancel()
+
+	// if a host dialer is specified, the dialer should be used to resolve the DNS
+	mock := NewMockTransportDialer(
+		[]string{},
+		[]string{},
+	)
+	connector.Dialer = MockHostTransportDialer{
+		Dialer: mock,
+		Host:   params.Host,
+	}
+
+	if mock.count != 0 {
+		t.Error("expecting no connections")
+	}
+	sqlDialer, _ := msdsn.ProtocolDialers["tcp"].(MssqlProtocolDialer)
+	conn, err := sqlDialer.DialSqlConnection(ctx, connector, &params)
+	if err != nil {
+		t.Error(err)
+	}
+
+	if mock.count != 1 {
+		t.Error("expecting 1 connection")
+	}
+
+	err = conn.Close()
+	if err != nil {
+		t.Error(err)
+	}
+
+	// if it is not a host dialer, the dialer should not be used to resolve DNS
+	connector.Dialer = mock
+	sqlDialer, _ = msdsn.ProtocolDialers["tcp"].(MssqlProtocolDialer)
+	_, err = sqlDialer.DialSqlConnection(ctx, connector, &params)
+	if err == nil {
+		t.Error(fmt.Errorf("dialer should not be used to resolve dns if not a host dialer"))
+	}
+}