Honor OAuth callback port during login

.nextchanges/cli/oauth-callback-port.md

@@ -0,0 +1 @@
+Fixed `databricks auth login` to honor `DATABRICKS_OAUTH_CALLBACK_PORT`.

cmd/auth/login.go

@@ -49,9 +49,22 @@ const (
 	// host used by the discovery login flow. Intended for testing and
 	// development against non-production environments. See WithDiscoveryHost
 	// in github.com/databricks/databricks-sdk-go/credentials/u2m.
-	discoveryHostEnvVar = "DATABRICKS_DISCOVERY_HOST"
+	discoveryHostEnvVar     = "DATABRICKS_DISCOVERY_HOST"
+	oauthCallbackPortEnvVar = "DATABRICKS_OAUTH_CALLBACK_PORT"
 )
 
+func withOAuthCallbackPort(ctx context.Context, opts []u2m.PersistentAuthOption) ([]u2m.PersistentAuthOption, error) {
+	value := env.Get(ctx, oauthCallbackPortEnvVar)
+	if value == "" {
+		return opts, nil
+	}
+	port, err := strconv.Atoi(value)
+	if err != nil {
+		return nil, fmt.Errorf("invalid %s value %q: %w", oauthCallbackPortEnvVar, value, err)
+	}
+	return append(opts, u2m.WithPort(port)), nil
+}
+
 // discoveryErr wraps an error (or creates a new one) and appends the
 // discovery fallback tip so users know they can bypass login.databricks.com.
 func discoveryErr(msg string, err error) error {
@@ -303,6 +316,10 @@ a new profile is created.
 			u2m.WithBrowser(getBrowserFunc(cmd)),
 			u2m.WithTokenCache(storage.WrapForOAuthArgument(ctx, tokenStore, mode, oauthArgument)),
 		}
+		persistentAuthOpts, err = withOAuthCallbackPort(ctx, persistentAuthOpts)
+		if err != nil {
+			return err
+		}
 		if len(scopesList) > 0 {
 			persistentAuthOpts = append(persistentAuthOpts, u2m.WithScopes(scopesList))
 		}
@@ -661,6 +678,10 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error {
 		u2m.WithDiscoveryLogin(),
 		u2m.WithTokenCache(storage.WrapForOAuthArgument(ctx, in.tokenStore, in.mode, arg)),
 	}
+	opts, err = withOAuthCallbackPort(ctx, opts)
+	if err != nil {
+		return discoveryErr("setting OAuth callback port", err)
+	}
 	if len(scopesList) > 0 {
 		opts = append(opts, u2m.WithScopes(scopesList))
 	}

cmd/auth/login_test.go

@@ -10,6 +10,7 @@ import (
 	"net/http/httptest"
 	"os"
 	"path/filepath"
+	"reflect"
 	"sync"
 	"testing"
 	"time"
@@ -80,12 +81,13 @@ func (f *fakeDiscoveryPersistentAuth) Close() error {
 }
 
 type fakeDiscoveryClient struct {
-	oauthArg          *u2m.BasicDiscoveryOAuthArgument
-	oauthArgErr       error
-	persistentAuth    discoveryPersistentAuth
-	persistentAuthErr error
-	introspection     *auth.IntrospectionResult
-	introspectionErr  error
+	oauthArg           *u2m.BasicDiscoveryOAuthArgument
+	oauthArgErr        error
+	persistentAuth     discoveryPersistentAuth
+	persistentAuthErr  error
+	persistentAuthOpts []u2m.PersistentAuthOption
+	introspection      *auth.IntrospectionResult
+	introspectionErr   error
 	// For assertions
 	introspectHost  string
 	introspectToken string
@@ -99,12 +101,42 @@ func (f *fakeDiscoveryClient) NewOAuthArgument(profileName string) (*u2m.BasicDi
 }
 
 func (f *fakeDiscoveryClient) NewPersistentAuth(ctx context.Context, opts ...u2m.PersistentAuthOption) (discoveryPersistentAuth, error) {
+	f.persistentAuthOpts = opts
 	if f.persistentAuthErr != nil {
 		return nil, f.persistentAuthErr
 	}
 	return f.persistentAuth, nil
 }
 
+func TestDiscoveryLogin_UsesOAuthCallbackPort(t *testing.T) {
+	t.Setenv("DATABRICKS_OAUTH_CALLBACK_PORT", "8030")
+	oauthArg, err := u2m.NewBasicDiscoveryOAuthArgument("DISCOVERY")
+	require.NoError(t, err)
+	dc := &fakeDiscoveryClient{
+		oauthArg: oauthArg,
+		persistentAuth: &fakeDiscoveryPersistentAuth{
+			challengeErr: errors.New("stop after creating persistent auth"),
+		},
+	}
+
+	ctx, _ := cmdio.NewTestContextWithStdout(t.Context())
+	err = discoveryLogin(ctx, discoveryLoginInputs{
+		dc:          dc,
+		profileName: "DISCOVERY",
+		timeout:     time.Second,
+		browserFunc: func(string) error { return nil },
+		tokenStore:  newTestStore(),
+	})
+	require.Error(t, err)
+
+	persistentAuth := &u2m.PersistentAuth{}
+	for _, opt := range dc.persistentAuthOpts {
+		opt(persistentAuth)
+	}
+	port := reflect.ValueOf(persistentAuth).Elem().FieldByName("port").Int()
+	assert.EqualValues(t, 8030, port)
+}
+
 func (f *fakeDiscoveryClient) IntrospectToken(ctx context.Context, host, accessToken string) (*auth.IntrospectionResult, error) {
 	f.introspectHost = host
 	f.introspectToken = accessToken