Update u2m port selection mechanism to start from 8020. (#1269)

## What changes are proposed in this pull request?

This PR replaces the port selection mechanism in `U2M` to start from
`8020`, and fallback incrementally (`8021`, `8022`, and so on... until
8040) if a port is not free. The PR also provides an environment/config
variable to specify a specific port for the OAuth callback. If set, that
port disables the fallback mechanism.

This change fixes a bug in the Databricks CLI which is only allowlisted
for port `8020` at the moment.

## How is this tested?

Unit tests

Author: renaudhartert-db

NEXT_CHANGELOG.md

@@ -6,6 +6,11 @@
 
 ### Bug Fixes
 
+- Update the `U2M` port selection mechanism to try port `8020` first and fall
+  back incrementally (to port `8021`, `8022`, and so on...) if a port is not 
+  available. This fixes an issue with the Databricks CLI which is only 
+  allowlisted on port `8020`. 
+
 ### Documentation
 
 ### Internal Changes

config/auth_u2m.go

@@ -44,7 +44,7 @@ func (u u2mCredentials) Configure(ctx context.Context, cfg *Config) (credentials
 	}
 
 	if u.ts == nil {
-		auth, err := u2m.NewPersistentAuth(ctx, u2m.WithOAuthArgument(arg))
+		auth, err := u2m.NewPersistentAuth(ctx, u2m.WithOAuthArgument(arg), u2m.WithPort(cfg.OAuthCallbackPort))
 		if err != nil {
 			logger.Debugf(ctx, "failed to create persistent auth: %v, continuing", err)
 			return nil, nil

config/config.go

@@ -70,6 +70,11 @@ type Config struct {
 	// in ~/.databrickscfg.
 	ConfigFile string `name:"config_file" env:"DATABRICKS_CONFIG_FILE"`
 
+	// OAuthPort is the port to use for the OAuth2 callback server. If not set,
+	// the default port with fallback is used. This means that setting a port
+	// will disable the fallback mechanism.
+	OAuthCallbackPort int `name:"oauth_callback_port" env:"DATABRICKS_OAUTH_CALLBACK_PORT" auth:"-"`
+
 	GoogleServiceAccount string `name:"google_service_account" env:"DATABRICKS_GOOGLE_SERVICE_ACCOUNT" auth:"google" auth_types:"google-id"`
 	GoogleCredentials    string `name:"google_credentials" env:"GOOGLE_CREDENTIALS" auth:"google,sensitive" auth_types:"google-credentials"`
 

credentials/u2m/persistent_auth.go

@@ -15,7 +15,6 @@ import (
 	cache "github.com/databricks/databricks-sdk-go/credentials/u2m/cache"
 	"github.com/databricks/databricks-sdk-go/httpclient"
 	"github.com/databricks/databricks-sdk-go/logger"
-	"github.com/databricks/databricks-sdk-go/retries"
 	"github.com/pkg/browser"
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/authhandler"
@@ -25,15 +24,25 @@ const (
 	// appClientId is the default client ID used by the SDK for U2M OAuth.
 	appClientID = "databricks-cli"
 
-	// appRedirectAddr is the default address for the OAuth2 callback server.
-	// Using ":0" tells the system to pick a random available port.
-	appRedirectAddr = "localhost:0"
+	// defaultPort is the default port for the OAuth2 callback server. If the
+	// port is already in use, the next port is tried (8021, 8022, etc.).
+	defaultPort = 8020
 
-	// listenerTimeout is the maximum amount of time to acquire listener on
-	// appRedirectAddr.
+	// maxPortFallback is the maximum port to try when using the fallback
+	// mechanism.
+	maxPortFallback = 8040
+
+	// listenerTimeout is the maximum duration spent trying to acquire a
+	// listener (including port selection).
 	listenerTimeout = 45 * time.Second
 )
 
+var (
+	// Internal errors used for testing.
+	errListenerTimeout = errors.New("failed to listen on any port: timeout")
+	errNoPortAvailable = errors.New("no port available to listen on")
+)
+
 // PersistentAuth is an OAuth manager that handles the U2M OAuth flow. Tokens
 // are stored in and looked up from the provided cache. Tokens include the
 // refresh token. On load, if the access token is expired, it is refreshed
@@ -44,25 +53,41 @@ const (
 type PersistentAuth struct {
 	// cache is the token cache to store and lookup tokens.
 	cache cache.TokenCache
+
 	// client is the HTTP client to use for OAuth2 requests.
 	client *http.Client
+
 	// endpointSupplier is the HTTP endpointSupplier to use for OAuth2 requests.
 	endpointSupplier OAuthEndpointSupplier
+
 	// oAuthArgument defines the workspace or account to authenticate to and the
 	// cache key for the token.
 	oAuthArgument OAuthArgument
+
 	// browser is the function to open a URL in the default browser.
 	browser func(url string) error
+
 	// ln is the listener for the OAuth2 callback server.
 	ln net.Listener
+
 	// ctx is the context to use for underlying operations. This is needed in
 	// order to implement the oauth2.TokenSource interface.
 	ctx context.Context
+
 	// redirectAddr is the redirect address for OAuth2 callbacks. The value is
 	// set to localhost:PORT by startListener which will dynamically assign a
 	// random port. If a value is already provided, it will be used instead
 	// (e.g. for testing).
 	redirectAddr string
+
+	// Optional port to use for the OAuth2 callback server. If set to 0, the
+	// default port with fallback is used. This means that setting a port will
+	// disable the fallback mechanism.
+	port int
+
+	// netListen is an optional function to listen on a TCP address. If not set,
+	// it will use net.Listen by default. This is useful for testing.
+	netListen func(network, address string) (net.Listener, error)
 }
 
 type PersistentAuthOption func(*PersistentAuth)
@@ -103,6 +128,13 @@ func WithBrowser(b func(url string) error) PersistentAuthOption {
 	}
 }
 
+// WithPort sets the port for the PersistentAuth.
+func WithPort(port int) PersistentAuthOption {
+	return func(a *PersistentAuth) {
+		a.port = port
+	}
+}
+
 // NewPersistentAuth creates a new PersistentAuth with the provided options.
 func NewPersistentAuth(ctx context.Context, opts ...PersistentAuthOption) (*PersistentAuth, error) {
 	p := &PersistentAuth{}
@@ -273,35 +305,51 @@ func (a *PersistentAuth) Challenge() error {
 
 // startListener starts a listener on appRedirectAddr, retrying if the address
 // is already in use.
-func (a *PersistentAuth) startListener(ctx context.Context) error {
-	// Use the value of redirectURL if it is already set. This is only expected
-	// in tests to set a fixed redirect URL.
-	addr := a.redirectAddr
-	if addr == "" {
-		addr = appRedirectAddr
-	}
-
-	listener, err := retries.Poll(ctx, listenerTimeout,
-		func() (*net.Listener, *retries.Err) {
-			var lc net.ListenConfig
-			l, err := lc.Listen(ctx, "tcp", addr)
-			if err != nil {
-				logger.Debugf(ctx, "failed to listen on %s: %v, retrying", addr, err)
-				return nil, retries.Continue(err)
-			}
-			return &l, nil
-		})
-	if err != nil {
-		return fmt.Errorf("listener: %w", err)
+func (pa *PersistentAuth) startListener(ctx context.Context) error {
+	if pa.port != 0 { // if port is set, use it
+		return pa.startListenerWithPort(pa.port)
 	}
-	a.ln = *listener
+	return pa.startListenerWithFallback(ctx)
+}
 
-	// Get the actual address that was assigned (including the port).
-	a.redirectAddr = a.ln.Addr().String()
-	logger.Debugf(ctx, "OAuth callback server listening on %s", a.redirectAddr)
+// startListenerWithFallback starts a listener that will try to find a free
+// port to listen on starting from the default port and incrementing by 1 until
+// a free port is found.
+func (pa *PersistentAuth) startListenerWithFallback(ctx context.Context) error {
+	startTime := time.Now()
+	for port := defaultPort; port <= maxPortFallback; port++ {
+		if time.Since(startTime) > listenerTimeout {
+			return errListenerTimeout
+		}
+		if err := pa.startListenerWithPort(port); err != nil {
+			logger.Debugf(ctx, "failed to listen on %d: %v, retrying", port, err)
+			continue
+		}
+		logger.Debugf(ctx, "OAuth callback server listening on %s", pa.redirectAddr)
+		return nil
+	}
+	return errNoPortAvailable
+}
+
+func (pa *PersistentAuth) startListenerWithPort(port int) error {
+	addr := fmt.Sprintf("localhost:%d", port)
+	listener, err := pa.listen("tcp", addr)
+	if err != nil {
+		return fmt.Errorf("failed to listen on %s: %w", addr, err)
+	}
+	pa.ln = listener
+	pa.redirectAddr = addr
 	return nil
 }
 
+func (pa *PersistentAuth) listen(network, addr string) (net.Listener, error) {
+	if pa.netListen != nil {
+		return pa.netListen(network, addr)
+	} else {
+		return net.Listen(network, addr)
+	}
+}
+
 func (a *PersistentAuth) Close() error {
 	if a.ln == nil {
 		return nil

credentials/u2m/persistent_auth_test.go

@@ -4,9 +4,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net"
 	"net/http"
 	"net/url"
-	"regexp"
 	"strings"
 	"testing"
 	"time"
@@ -298,9 +298,6 @@ func TestChallenge(t *testing.T) {
 	}
 	defer p.Close()
 
-	// Set a fixed redirect URL for the test.
-	p.redirectAddr = "localhost:1337"
-
 	errc := make(chan error)
 	go func() {
 		err := p.Challenge()
@@ -309,7 +306,7 @@ func TestChallenge(t *testing.T) {
 	}()
 
 	state := <-browserOpened
-	resp, err := http.Get(fmt.Sprintf("http://localhost:1337?code=__THIS__&state=%s", state))
+	resp, err := http.Get(fmt.Sprintf("http://localhost:8020?code=__THIS__&state=%s", state))
 	if err != nil {
 		t.Fatalf("http.Get(): want no error, got %v", err)
 	}
@@ -351,8 +348,6 @@ func TestChallenge_ReturnsErrorOnFailure(t *testing.T) {
 	}
 	defer p.Close()
 
-	p.redirectAddr = "localhost:1337" // set a fixed redirect URL for the test
-
 	errc := make(chan error)
 	go func() {
 		err := p.Challenge()
@@ -361,7 +356,7 @@ func TestChallenge_ReturnsErrorOnFailure(t *testing.T) {
 	}()
 
 	<-browserOpened
-	resp, err := http.Get("http://localhost:1337?error=access_denied&error_description=Policy%20evaluation%20failed%20for%20this%20request")
+	resp, err := http.Get("http://localhost:8020?error=access_denied&error_description=Policy%20evaluation%20failed%20for%20this%20request")
 	if err != nil {
 		t.Fatalf("http.Get(): want no error, got %v", err)
 	}
@@ -380,36 +375,108 @@ func TestChallenge_ReturnsErrorOnFailure(t *testing.T) {
 	}
 }
 
-// Verifies that startListener assigns a random port to the redirectAddr.
-func TestPersistentAuth_startListener_useDifferentPorts(t *testing.T) {
-	ctx := context.Background()
-	arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.com", "xyz")
-	if err != nil {
-		t.Fatalf("NewBasicAccountOAuthArgument(): want no error, got %v", err)
+func TestPersistentAuth_startListener_startFrom8020(t *testing.T) {
+	pa := &PersistentAuth{}
+	pa.netListen = func(_, address string) (net.Listener, error) {
+		return nil, nil
 	}
 
-	p1, err := NewPersistentAuth(ctx, WithOAuthArgument(arg))
-	if err != nil {
-		t.Fatalf("NewPersistentAuth(): want no error, got %v", err)
+	gotErr := pa.startListener(context.Background())
+
+	if gotErr != nil {
+		t.Fatalf("pa.startListener(): want no error, got %v", gotErr)
 	}
-	defer p1.Close()
+	if pa.redirectAddr != "localhost:8020" {
+		t.Errorf("pa.redirectAddr should be localhost:8020, got %s", pa.redirectAddr)
+	}
+}
 
-	p2, err := NewPersistentAuth(ctx, WithOAuthArgument(arg))
-	if err != nil {
-		t.Fatalf("NewPersistentAuth(): want no error, got %v", err)
+func TestPersistentAuth_startListener_incrementalFallBack(t *testing.T) {
+	pa := &PersistentAuth{}
+	pa.netListen = func(_, address string) (net.Listener, error) {
+		if address == "localhost:8020" {
+			return nil, fmt.Errorf("address already in use")
+		}
+		if address == "localhost:8021" {
+			return nil, fmt.Errorf("address already in use")
+		}
+		return nil, nil
+	}
+
+	gotErr := pa.startListener(context.Background())
+
+	if gotErr != nil {
+		t.Fatalf("pa.startListener(): want no error, got %v", gotErr)
+	}
+	if pa.redirectAddr != "localhost:8022" {
+		t.Errorf("pa.redirectAddr should be localhost:8022, got %s", pa.redirectAddr)
+	}
+}
+
+func TestPersistentAuth_startListener_noAvailablePort(t *testing.T) {
+	pa := &PersistentAuth{}
+	pa.netListen = func(_, address string) (net.Listener, error) {
+		return nil, fmt.Errorf("address already in use")
 	}
-	defer p2.Close()
 
-	p1.startListener(ctx)
-	p2.startListener(ctx)
+	gotErr := pa.startListener(context.Background())
+
+	if !errors.Is(gotErr, errNoPortAvailable) {
+		t.Fatalf("pa.startListener(): want error %v, got %v", errNoPortAvailable, gotErr)
+	}
+}
 
-	if !regexp.MustCompile(`^127\.0\.0\.1:\d+$`).MatchString(p1.redirectAddr) {
-		t.Errorf("p1.redirectAddr should be random localhost port, got %s", p1.redirectAddr)
+func TestPersistentAuth_startListener_maxPortFallbackIncluded(t *testing.T) {
+	maxAddress := fmt.Sprintf("localhost:%d", maxPortFallback)
+	pa := &PersistentAuth{}
+	pa.netListen = func(_, address string) (net.Listener, error) {
+		if address == maxAddress {
+			return nil, nil
+		}
+		return nil, fmt.Errorf("address already in use")
 	}
-	if !regexp.MustCompile(`^127\.0\.0\.1:\d+$`).MatchString(p2.redirectAddr) {
-		t.Errorf("p2.redirectAddr should be random localhost port, got %s", p2.redirectAddr)
+
+	gotErr := pa.startListener(context.Background())
+
+	if gotErr != nil {
+		t.Fatalf("pa.startListener(): want no error, got %v", gotErr)
+	}
+	if pa.redirectAddr != maxAddress {
+		t.Errorf("pa.redirectAddr should be %s, got %s", maxAddress, pa.redirectAddr)
+	}
+}
+
+func TestPersistentAuth_startListener_explicitPort(t *testing.T) {
+	explicitPort := 1337
+	pa := &PersistentAuth{port: explicitPort}
+	pa.netListen = func(_, address string) (net.Listener, error) {
+		return nil, nil
 	}
-	if p1.redirectAddr == p2.redirectAddr {
-		t.Errorf("p1.redirectURL and p2.redirectURL should be different, got %s", p1.redirectAddr)
+
+	gotErr := pa.startListener(context.Background())
+
+	if gotErr != nil {
+		t.Fatalf("pa.startListener(): want no error, got %v", gotErr)
+	}
+	if pa.redirectAddr != "localhost:1337" {
+		t.Errorf("pa.redirectAddr should be localhost:1337, got %s", pa.redirectAddr)
+	}
+}
+
+func TestPersistentAuth_startListener_explicitPortNoFallBack(t *testing.T) {
+	testError := errors.New("test error")
+	explicitPort := 1337
+	pa := &PersistentAuth{port: explicitPort}
+	pa.netListen = func(_, address string) (net.Listener, error) {
+		if address == "localhost:1337" {
+			return nil, testError
+		}
+		return nil, nil
+	}
+
+	gotErr := pa.startListener(context.Background())
+
+	if !errors.Is(gotErr, testError) {
+		t.Fatalf("pa.startListener(): want error %v, got %v", testError, gotErr)
 	}
 }