Allow for local dns resolution with a custom dialer

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 !strings.Contains(err.Error(), "no such host") {
+		t.Error(fmt.Errorf("dialer should not be used to resolve dns if not a host dialer"))
+	}
+}