Repository: supabase/cli
Pull request #900 (2023-03-02 07:38:37)

chore: refactor dns resolver for better test coverage
## What kind of change does this PR introduce?

unit tests

## What is the new behavior?

- allow ip lookup with host name only
- improve test coverage

## Additional context

Add any other context or screenshots.


coveralls: ## Pull Request Test Coverage Report for [Build 4312134561](https://coveralls.io/builds/57458062)

* **28** of **29**   **(96.55%)**  changed or added relevant lines in **2** files are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage increased (+**0.2%**) to **62.575%**

---

|  Changes Missing Coverage | Covered Lines | Changed/Added Lines | % |
| :-----|--------------|--------|---: |
| [internal/utils/connect.go](https://coveralls.io/builds/57458062/source?filename=internal%2Futils%2Fconnect.go#L34) | 0 | 1 | 0.0%
<!-- | **Total:** | **28** | **29** | **96.55%** | -->


|  Totals | [![Coverage Status](https://coveralls.io/builds/57458062/badge)](https://coveralls.io/builds/57458062) |
| :-- | --: |
| Change from base [Build 4310500187](https://coveralls.io/builds/57451622): |  0.2% |
| Covered Lines: | 3757 |
| Relevant Lines: | 6004 |

---
##### 💛  - [Coveralls](https://coveralls.io)


Patch:
diff --git a/go.mod b/go.mod
index a2ca0eee..43072f39 100644
--- a/go.mod
+++ b/go.mod
@@ -64,6 +64,7 @@ require (
 	github.com/pelletier/go-toml/v2 v2.0.6 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/sahilm/fuzzy v0.1.0 // indirect
+	github.com/stretchr/objx v0.5.0 // indirect
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
 	github.com/ugorji/go/codec v1.2.9 // indirect
 	github.com/valyala/bytebufferpool v1.0.0 // indirect
diff --git a/internal/utils/api.go b/internal/utils/api.go
index 52ffad4d..bd9526f5 100644
--- a/internal/utils/api.go
+++ b/internal/utils/api.go
@@ -53,13 +53,9 @@ type dnsResponse struct {
 }
 
 // Performs DNS lookup via HTTPS, in case firewall blocks native netgo resolver.
-func FallbackLookupIP(ctx context.Context, address string) ([]string, error) {
-	host, port, err := net.SplitHostPort(address)
-	if err != nil {
-		return nil, err
-	}
+func FallbackLookupIP(ctx context.Context, host string) ([]string, error) {
 	if net.ParseIP(host) != nil {
-		return []string{address}, nil
+		return []string{host}, nil
 	}
 	// Ref: https://developers.cloudflare.com/1.1.1.1/encryption/dns-over-https/make-api-requests/dns-json
 	req, err := http.NewRequestWithContext(ctx, "GET", "https://1.1.1.1/dns-query?name="+host, nil)
@@ -90,11 +86,11 @@ func FallbackLookupIP(ctx context.Context, address string) ([]string, error) {
 	var resolved []string
 	for _, answer := range data.Answer {
 		if answer.Type == dnsIPv4Type || answer.Type == dnsIPv6Type {
-			resolved = append(resolved, net.JoinHostPort(answer.Data, port))
+			resolved = append(resolved, answer.Data)
 		}
 	}
 	if len(resolved) == 0 {
-		err = fmt.Errorf("failed to locate valid IP for %s; resolves to %#v", address, data.Answer)
+		err = fmt.Errorf("failed to locate valid IP for %s; resolves to %#v", host, data.Answer)
 	}
 	return resolved, err
 }
@@ -188,6 +184,37 @@ func WithTraceContext(ctx context.Context) context.Context {
 	return httptrace.WithClientTrace(ctx, trace)
 }
 
+type DialContextFunc func(context.Context, string, string) (net.Conn, error)
+
+// Wraps a DialContext with DNS-over-HTTPS as fallback resolver
+func withFallbackDNS(dialContext DialContextFunc) DialContextFunc {
+	dnsOverHttps := func(ctx context.Context, network, address string) (net.Conn, error) {
+		host, port, err := net.SplitHostPort(address)
+		if err != nil {
+			return nil, err
+		}
+		ip, err := FallbackLookupIP(ctx, host)
+		if err != nil {
+			return nil, err
+		}
+		return dialContext(ctx, network, net.JoinHostPort(ip[0], port))
+	}
+	if DNSResolver.Value == DNS_OVER_HTTPS {
+		return dnsOverHttps
+	}
+	nativeWithFallback := func(ctx context.Context, network, address string) (net.Conn, error) {
+		conn, err := dialContext(ctx, network, address)
+		// Workaround when pure Go DNS resolver fails https://github.com/golang/go/issues/12524
+		if err, ok := err.(net.Error); ok && err.Timeout() {
+			if conn, err := dnsOverHttps(ctx, network, address); err == nil {
+				return conn, err
+			}
+		}
+		return conn, err
+	}
+	return nativeWithFallback
+}
+
 func GetSupabase() *supabase.ClientWithResponses {
 	clientOnce.Do(func() {
 		token, err := LoadAccessToken()
@@ -199,24 +226,7 @@ func GetSupabase() *supabase.ClientWithResponses {
 			log.Fatalln(err)
 		}
 		if t, ok := http.DefaultTransport.(*http.Transport); ok {
-			dialContext := t.DialContext
-			t.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
-				if DNSResolver.Value == DNS_OVER_HTTPS {
-					ip, err := FallbackLookupIP(ctx, address)
-					if err != nil {
-						return nil, err
-					}
-					return dialContext(ctx, network, ip[0])
-				}
-				conn, err := dialContext(ctx, network, address)
-				// Workaround when pure Go DNS resolver fails https://github.com/golang/go/issues/12524
-				if err, ok := err.(net.Error); ok && err.Timeout() {
-					if ip, err := FallbackLookupIP(ctx, address); err == nil {
-						return dialContext(ctx, network, ip[0])
-					}
-				}
-				return conn, err
-			}
+			t.DialContext = withFallbackDNS(t.DialContext)
 		}
 		apiClient, err = supabase.NewClientWithResponses(
 			GetSupabaseAPIHost(),
diff --git a/internal/utils/connect.go b/internal/utils/connect.go
index 4fb1f6f9..dea82aac 100644
--- a/internal/utils/connect.go
+++ b/internal/utils/connect.go
@@ -3,10 +3,8 @@ package utils
 import (
 	"context"
 	"fmt"
-	"net"
 	"net/url"
 	"os"
-	"strconv"
 
 	"github.com/jackc/pgconn"
 	"github.com/jackc/pgx/v4"
@@ -33,10 +31,7 @@ func ConnectRemotePostgres(ctx context.Context, username, password, database, ho
 	opts := append(options, func(cc *pgx.ConnConfig) {
 		cc.PreferSimpleProtocol = true
 		if DNSResolver.Value == DNS_OVER_HTTPS {
-			cc.LookupFunc = func(ctx context.Context, host string) ([]string, error) {
-				address := net.JoinHostPort(host, strconv.FormatUint(uint64(cc.Port), 10))
-				return FallbackLookupIP(ctx, address)
-			}
+			cc.LookupFunc = FallbackLookupIP
 		}
 	})
 	conn, err := ConnectByUrl(ctx, pgUrl, opts...)


Test patch:
diff --git a/internal/utils/api_test.go b/internal/utils/api_test.go
index 3f110641..742e3333 100644
--- a/internal/utils/api_test.go
+++ b/internal/utils/api_test.go
@@ -3,10 +3,12 @@ package utils
 import (
 	"context"
 	"errors"
+	"net"
 	"net/http"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/supabase/cli/internal/testing/apitest"
 	"gopkg.in/h2non/gock.v1"
 )
@@ -26,10 +28,10 @@ func TestLookupIP(t *testing.T) {
 				{Type: dnsIPv4Type, Data: "127.0.0.1"},
 			}})
 		// Run test
-		ip, err := FallbackLookupIP(context.Background(), host+":443")
+		ip, err := FallbackLookupIP(context.Background(), host)
 		// Validate output
 		assert.NoError(t, err)
-		assert.ElementsMatch(t, []string{"127.0.0.1:443"}, ip)
+		assert.ElementsMatch(t, []string{"127.0.0.1"}, ip)
 		assert.Empty(t, apitest.ListUnmatchedRequests())
 	})
 
@@ -46,28 +48,22 @@ func TestLookupIP(t *testing.T) {
 				{Type: dnsIPv6Type, Data: "2606:2800:220:1:248:1893:25c8:1946"},
 			}})
 		// Run test
-		ip, err := FallbackLookupIP(context.Background(), "api.supabase.com:443")
+		ip, err := FallbackLookupIP(context.Background(), "api.supabase.com")
 		// Validate output
 		assert.NoError(t, err)
-		assert.ElementsMatch(t, []string{"[2606:2800:220:1:248:1893:25c8:1946]:443"}, ip)
+		assert.ElementsMatch(t, []string{"2606:2800:220:1:248:1893:25c8:1946"}, ip)
 		assert.Empty(t, apitest.ListUnmatchedRequests())
 	})
 
 	t.Run("returns immediately if already resolved", func(t *testing.T) {
 		// Run test
-		ip, err := FallbackLookupIP(context.Background(), "127.0.0.1:443")
+		ip, err := FallbackLookupIP(context.Background(), "127.0.0.1")
 		// Validate output
 		assert.NoError(t, err)
-		assert.ElementsMatch(t, []string{"127.0.0.1:443"}, ip)
+		assert.ElementsMatch(t, []string{"127.0.0.1"}, ip)
 		assert.Empty(t, apitest.ListUnmatchedRequests())
 	})
 
-	t.Run("empty on malformed address", func(t *testing.T) {
-		ip, err := FallbackLookupIP(context.Background(), "bad?url")
-		assert.ErrorContains(t, err, "missing port in address")
-		assert.Empty(t, ip)
-	})
-
 	t.Run("empty on network failure", func(t *testing.T) {
 		// Setup http mock
 		defer gock.OffAll()
@@ -77,7 +73,7 @@ func TestLookupIP(t *testing.T) {
 			MatchHeader("accept", "application/dns-json").
 			ReplyError(errors.New("network error"))
 		// Run test
-		ip, err := FallbackLookupIP(context.Background(), host+":443")
+		ip, err := FallbackLookupIP(context.Background(), host)
 		// Validate output
 		assert.ErrorContains(t, err, "network error")
 		assert.Empty(t, ip)
@@ -93,7 +89,7 @@ func TestLookupIP(t *testing.T) {
 			MatchHeader("accept", "application/dns-json").
 			Reply(http.StatusServiceUnavailable)
 		// Run test
-		ip, err := FallbackLookupIP(context.Background(), host+":443")
+		ip, err := FallbackLookupIP(context.Background(), host)
 		// Validate output
 		assert.ErrorContains(t, err, "status 503")
 		assert.Empty(t, ip)
@@ -110,7 +106,7 @@ func TestLookupIP(t *testing.T) {
 			Reply(http.StatusOK).
 			JSON("malformed")
 		// Run test
-		ip, err := FallbackLookupIP(context.Background(), host+":443")
+		ip, err := FallbackLookupIP(context.Background(), host)
 		// Validate output
 		assert.ErrorContains(t, err, "invalid character 'm' looking for beginning of value")
 		assert.Empty(t, ip)
@@ -127,9 +123,9 @@ func TestLookupIP(t *testing.T) {
 			Reply(http.StatusOK).
 			JSON(&dnsResponse{})
 		// Run test
-		ip, err := FallbackLookupIP(context.Background(), host+":443")
+		ip, err := FallbackLookupIP(context.Background(), host)
 		// Validate output
-		assert.ErrorContains(t, err, "failed to locate valid IP for api.supabase.io:443; resolves to []utils.dnsAnswer(nil)")
+		assert.ErrorContains(t, err, "failed to locate valid IP for api.supabase.io; resolves to []utils.dnsAnswer(nil)")
 		assert.Empty(t, ip)
 		assert.Empty(t, apitest.ListUnmatchedRequests())
 	})
@@ -190,5 +186,113 @@ func TestResolveCNAME(t *testing.T) {
 		assert.ErrorContains(t, err, "failed to locate appropriate CNAME record for api.supabase.io")
 		assert.Empty(t, apitest.ListUnmatchedRequests())
 	})
+}
+
+type MockDialer struct {
+	mock.Mock
+}
 
+func (m *MockDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
+	args := m.Called(ctx, network, address)
+	if conn, ok := args.Get(0).(net.Conn); ok {
+		return conn, args.Error(1)
+	}
+	return nil, args.Error(1)
+}
+
+func TestFallbackDNS(t *testing.T) {
+	errNetwork := errors.New("network error")
+	errDNS := &net.DNSError{
+		IsTimeout: true,
+	}
+
+	t.Run("overrides DialContext with DoH", func(t *testing.T) {
+		DNSResolver.Value = DNS_OVER_HTTPS
+		// Setup mock dialer
+		dialer := MockDialer{}
+		dialer.On("DialContext", mock.Anything, mock.Anything, "127.0.0.1:80").
+			Return(nil, errNetwork)
+		wrapped := withFallbackDNS(dialer.DialContext)
+		// Setup http mock
+		defer gock.OffAll()
+		gock.New("https://1.1.1.1").
+			Get("/dns-query").
+			MatchParam("name", host).
+			MatchHeader("accept", "application/dns-json").
+			Reply(http.StatusOK).
+			JSON(&dnsResponse{Answer: []dnsAnswer{
+				{Type: dnsIPv4Type, Data: "127.0.0.1"},
+			}})
+		// Run test
+		conn, err := wrapped(context.Background(), "udp", host+":80")
+		// Check error
+		assert.ErrorIs(t, err, errNetwork)
+		assert.Nil(t, conn)
+		dialer.AssertExpectations(t)
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+
+	t.Run("native with DoH fallback", func(t *testing.T) {
+		DNSResolver.Value = DNS_GO_NATIVE
+		// Setup mock dialer
+		dialer := MockDialer{}
+		dialer.On("DialContext", mock.Anything, mock.Anything, host+":80").
+			Return(nil, errDNS)
+		dialer.On("DialContext", mock.Anything, mock.Anything, "127.0.0.1:80").
+			Return(nil, nil)
+		wrapped := withFallbackDNS(dialer.DialContext)
+		// Setup http mock
+		defer gock.OffAll()
+		gock.New("https://1.1.1.1").
+			Get("/dns-query").
+			MatchParam("name", host).
+			MatchHeader("accept", "application/dns-json").
+			Reply(http.StatusOK).
+			JSON(&dnsResponse{Answer: []dnsAnswer{
+				{Type: dnsIPv4Type, Data: "127.0.0.1"},
+			}})
+		// Run test
+		conn, err := wrapped(context.Background(), "udp", host+":80")
+		// Check error
+		assert.NoError(t, err)
+		assert.Nil(t, conn)
+		dialer.AssertExpectations(t)
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+
+	t.Run("throws error on malformed address", func(t *testing.T) {
+		DNSResolver.Value = DNS_OVER_HTTPS
+		// Setup mock dialer
+		dialer := MockDialer{}
+		wrapped := withFallbackDNS(dialer.DialContext)
+		// Run test
+		conn, err := wrapped(context.Background(), "udp", "bad?url")
+		// Check error
+		assert.ErrorContains(t, err, "missing port in address")
+		assert.Nil(t, conn)
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
+
+	t.Run("throws error on fallback failure", func(t *testing.T) {
+		DNSResolver.Value = DNS_GO_NATIVE
+		// Setup mock dialer
+		dialer := MockDialer{}
+		dialer.On("DialContext", mock.Anything, mock.Anything, host+":80").
+			Return(nil, errDNS)
+		wrapped := withFallbackDNS(dialer.DialContext)
+		// Setup http mock
+		defer gock.OffAll()
+		gock.New("https://1.1.1.1").
+			Get("/dns-query").
+			MatchParam("name", host).
+			MatchHeader("accept", "application/dns-json").
+			ReplyError(errNetwork)
+		// Run test
+		conn, err := wrapped(context.Background(), "udp", host+":80")
+		// Check error
+		assert.ErrorIs(t, err, errDNS)
+		assert.Nil(t, conn)
+		dialer.AssertExpectations(t)
+		assert.Empty(t, apitest.ListUnmatchedRequests())
+	})
 }