Repository: supabase/cli
Pull request #898 (2023-03-01 11:19:45)

feat: support overriding dns resolver with flag
## What kind of change does this PR introduce?

feature

## What is the new behavior?

- `supabase link --dns-resolver=https` will use dns-over-https resolver for both api and db hostname

## Additional context

Add any other context or screenshots.



Patch:
diff --git a/cmd/root.go b/cmd/root.go
index 5e67d71c..d07db1dc 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -80,6 +80,7 @@ func init() {
 	flags.Bool("debug", false, "output debug logs to stderr")
 	flags.String("workdir", "", "path to a Supabase project directory")
 	flags.Bool("experimental", false, "enable experimental features")
+	flags.Var(&utils.DNSResolver, "dns-resolver", "lookup domain names using the specified resolver")
 	cobra.CheckErr(viper.BindPFlags(flags))
 
 	rootCmd.SetVersionTemplate("{{.Version}}\n")
diff --git a/internal/utils/api.go b/internal/utils/api.go
index b91611bf..52ffad4d 100644
--- a/internal/utils/api.go
+++ b/internal/utils/api.go
@@ -4,7 +4,9 @@ import (
 	"context"
 	"crypto/tls"
 	"encoding/json"
+	"errors"
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"net/http"
@@ -18,10 +20,20 @@ import (
 	supabase "github.com/supabase/cli/pkg/api"
 )
 
+const (
+	DNS_GO_NATIVE  = "native"
+	DNS_OVER_HTTPS = "https"
+)
+
 var (
 	clientOnce sync.Once
 	apiClient  *supabase.ClientWithResponses
 	httpClient = http.Client{Timeout: 10 * time.Second}
+
+	DNSResolver = EnumFlag{
+		Allowed: []string{DNS_GO_NATIVE, DNS_OVER_HTTPS},
+		Value:   DNS_GO_NATIVE,
+	}
 )
 
 const (
@@ -41,39 +53,50 @@ type dnsResponse struct {
 }
 
 // Performs DNS lookup via HTTPS, in case firewall blocks native netgo resolver.
-func fallbackLookupIP(ctx context.Context, address string) string {
+func FallbackLookupIP(ctx context.Context, address string) ([]string, error) {
 	host, port, err := net.SplitHostPort(address)
 	if err != nil {
-		return ""
+		return nil, err
+	}
+	if net.ParseIP(host) != nil {
+		return []string{address}, 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)
 	if err != nil {
-		return ""
+		return nil, err
 	}
 	req.Header.Add("accept", "application/dns-json")
 	// Sends request
 	resp, err := httpClient.Do(req)
 	if err != nil {
-		return ""
+		return nil, err
 	}
 	defer resp.Body.Close()
 	if resp.StatusCode != http.StatusOK {
-		return ""
+		body, err := io.ReadAll(resp.Body)
+		if err != nil || len(body) == 0 {
+			body = []byte(fmt.Sprintf("status %d", resp.StatusCode))
+		}
+		return nil, errors.New(string(body))
 	}
 	// Parses response
 	var data dnsResponse
 	dec := json.NewDecoder(resp.Body)
 	if err := dec.Decode(&data); err != nil {
-		return ""
+		return nil, err
 	}
 	// Look for first valid IP
+	var resolved []string
 	for _, answer := range data.Answer {
 		if answer.Type == dnsIPv4Type || answer.Type == dnsIPv6Type {
-			return net.JoinHostPort(answer.Data, port)
+			resolved = append(resolved, net.JoinHostPort(answer.Data, port))
 		}
 	}
-	return ""
+	if len(resolved) == 0 {
+		err = fmt.Errorf("failed to locate valid IP for %s; resolves to %#v", address, data.Answer)
+	}
+	return resolved, err
 }
 
 func ResolveCNAME(ctx context.Context, host string) (string, error) {
@@ -178,11 +201,18 @@ func GetSupabase() *supabase.ClientWithResponses {
 		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 := fallbackLookupIP(ctx, address); ip != "" {
-						return dialContext(ctx, network, ip)
+					if ip, err := FallbackLookupIP(ctx, address); err == nil {
+						return dialContext(ctx, network, ip[0])
 					}
 				}
 				return conn, err
diff --git a/internal/utils/connect.go b/internal/utils/connect.go
index fc173e75..4fb1f6f9 100644
--- a/internal/utils/connect.go
+++ b/internal/utils/connect.go
@@ -3,8 +3,10 @@ package utils
 import (
 	"context"
 	"fmt"
+	"net"
 	"net/url"
 	"os"
+	"strconv"
 
 	"github.com/jackc/pgconn"
 	"github.com/jackc/pgx/v4"
@@ -30,6 +32,12 @@ func ConnectRemotePostgres(ctx context.Context, username, password, database, ho
 	//      don't benefit from per connection server side cache.
 	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)
+			}
+		}
 	})
 	conn, err := ConnectByUrl(ctx, pgUrl, opts...)
 	if !pgconn.Timeout(err) {


Test patch:
diff --git a/internal/utils/api_test.go b/internal/utils/api_test.go
index cce0f332..3f110641 100644
--- a/internal/utils/api_test.go
+++ b/internal/utils/api_test.go
@@ -11,9 +11,9 @@ import (
 	"gopkg.in/h2non/gock.v1"
 )
 
-func TestFallbackDNS(t *testing.T) {
-	const host = "api.supabase.io"
+const host = "api.supabase.io"
 
+func TestLookupIP(t *testing.T) {
 	t.Run("resolves IPv4 with CloudFlare", func(t *testing.T) {
 		// Setup http mock
 		defer gock.OffAll()
@@ -26,9 +26,10 @@ func TestFallbackDNS(t *testing.T) {
 				{Type: dnsIPv4Type, Data: "127.0.0.1"},
 			}})
 		// Run test
-		ip := fallbackLookupIP(context.Background(), host+":443")
+		ip, err := FallbackLookupIP(context.Background(), host+":443")
 		// Validate output
-		assert.Equal(t, "127.0.0.1:443", ip)
+		assert.NoError(t, err)
+		assert.ElementsMatch(t, []string{"127.0.0.1:443"}, ip)
 		assert.Empty(t, apitest.ListUnmatchedRequests())
 	})
 
@@ -45,14 +46,26 @@ func TestFallbackDNS(t *testing.T) {
 				{Type: dnsIPv6Type, Data: "2606:2800:220:1:248:1893:25c8:1946"},
 			}})
 		// Run test
-		ip := fallbackLookupIP(context.Background(), "api.supabase.com:443")
+		ip, err := FallbackLookupIP(context.Background(), "api.supabase.com:443")
 		// Validate output
-		assert.Equal(t, "[2606:2800:220:1:248:1893:25c8:1946]:443", ip)
+		assert.NoError(t, err)
+		assert.ElementsMatch(t, []string{"[2606:2800:220:1:248:1893:25c8:1946]:443"}, 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")
+		// Validate output
+		assert.NoError(t, err)
+		assert.ElementsMatch(t, []string{"127.0.0.1:443"}, ip)
 		assert.Empty(t, apitest.ListUnmatchedRequests())
 	})
 
 	t.Run("empty on malformed address", func(t *testing.T) {
-		assert.Equal(t, "", fallbackLookupIP(context.Background(), "bad?url"))
+		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) {
@@ -64,9 +77,10 @@ func TestFallbackDNS(t *testing.T) {
 			MatchHeader("accept", "application/dns-json").
 			ReplyError(errors.New("network error"))
 		// Run test
-		ip := fallbackLookupIP(context.Background(), host+":443")
+		ip, err := FallbackLookupIP(context.Background(), host+":443")
 		// Validate output
-		assert.Equal(t, "", ip)
+		assert.ErrorContains(t, err, "network error")
+		assert.Empty(t, ip)
 		assert.Empty(t, apitest.ListUnmatchedRequests())
 	})
 
@@ -79,9 +93,10 @@ func TestFallbackDNS(t *testing.T) {
 			MatchHeader("accept", "application/dns-json").
 			Reply(http.StatusServiceUnavailable)
 		// Run test
-		ip := fallbackLookupIP(context.Background(), host+":443")
+		ip, err := FallbackLookupIP(context.Background(), host+":443")
 		// Validate output
-		assert.Equal(t, "", ip)
+		assert.ErrorContains(t, err, "status 503")
+		assert.Empty(t, ip)
 		assert.Empty(t, apitest.ListUnmatchedRequests())
 	})
 
@@ -95,9 +110,10 @@ func TestFallbackDNS(t *testing.T) {
 			Reply(http.StatusOK).
 			JSON("malformed")
 		// Run test
-		ip := fallbackLookupIP(context.Background(), host+":443")
+		ip, err := FallbackLookupIP(context.Background(), host+":443")
 		// Validate output
-		assert.Equal(t, "", ip)
+		assert.ErrorContains(t, err, "invalid character 'm' looking for beginning of value")
+		assert.Empty(t, ip)
 		assert.Empty(t, apitest.ListUnmatchedRequests())
 	})
 
@@ -111,12 +127,15 @@ func TestFallbackDNS(t *testing.T) {
 			Reply(http.StatusOK).
 			JSON(&dnsResponse{})
 		// Run test
-		ip := fallbackLookupIP(context.Background(), host+":443")
+		ip, err := FallbackLookupIP(context.Background(), host+":443")
 		// Validate output
-		assert.Equal(t, "", ip)
+		assert.ErrorContains(t, err, "failed to locate valid IP for api.supabase.io:443; resolves to []utils.dnsAnswer(nil)")
+		assert.Empty(t, ip)
 		assert.Empty(t, apitest.ListUnmatchedRequests())
 	})
+}
 
+func TestResolveCNAME(t *testing.T) {
 	t.Run("resolves CNAMEs with CloudFlare", func(t *testing.T) {
 		defer gock.OffAll()
 		gock.New("https://1.1.1.1").