fix(auth): thread connector transport and User-Agent into OAuth M2M token requests

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Release History
 
 ## Unreleased
+- OAuth M2M token requests and the OIDC discovery lookup now go through the connector's transport (`WithTransport`) and carry the driver User-Agent, including `WithUserAgentEntry`, so they are attributable in `system.access.audit` instead of appearing as `Go-http-client` (databricks/databricks-sql-go#413).
 - **Behavior change:** `m2m.NewAuthenticatorWithScopes` now requests exactly the given scopes instead of always appending `all-apis`, so service principals with scoped OAuth secrets (e.g. `sql` only) can authenticate (databricks/databricks-sql-go#476). Empty scopes and `m2m.NewAuthenticator` still request `all-apis`; callers that relied on it being appended should pass no scopes, or add `all-apis` to their list if they need both. The SEA/kernel backend now forwards custom M2M scopes instead of rejecting them.
 
 ## v1.16.0 (2026-09-24)

auth/oauth/m2m/m2m.go

@@ -34,10 +34,22 @@ type authClient struct {
 	clientSecret string
 	hostName     string
 	scopes       []string
+	httpClient   *http.Client
 	tokenSource  oauth2.TokenSource
 	mx           sync.Mutex
 }
 
+// SetHTTPClient routes OIDC discovery and token requests through the given client
+// instead of http.DefaultClient, so they carry the connector's transport and
+// User-Agent (WithTransport, WithUserAgentEntry) like the rest of the driver's
+// traffic. The driver calls it structurally once the connector config is final;
+// it is a no-op on token sources already created.
+func (c *authClient) SetHTTPClient(client *http.Client) {
+	c.mx.Lock()
+	defer c.mx.Unlock()
+	c.httpClient = client
+}
+
 // M2MCredentials exposes the raw client-credentials so the SEA-via-kernel backend
 // can drive the kernel's own M2M flow, keeping cfg.Authenticator the single source
 // of truth for auth mode. It structurally satisfies the M2MCredentialsProvider
@@ -66,12 +78,16 @@ func (c *authClient) Authenticate(r *http.Request) error {
 		return nil
 	}
 
-	config, err := GetConfig(context.Background(), c.hostName, c.clientID, c.clientSecret, c.scopes)
+	ctx := context.Background()
+	if c.httpClient != nil {
+		ctx = context.WithValue(ctx, oauth2.HTTPClient, c.httpClient)
+	}
+	config, err := GetConfig(ctx, c.hostName, c.clientID, c.clientSecret, c.scopes)
 	if err != nil {
 		return fmt.Errorf("unable to generate clientCredentials.Config: %w", err)
 	}
 
-	c.tokenSource = GetTokenSource(config)
+	c.tokenSource = config.TokenSource(ctx)
 	token, err := c.tokenSource.Token()
 	if err != nil {
 		logger.Err(err).Msg("failed to get token")

auth/oauth/m2m/m2m_test.go

@@ -1,9 +1,13 @@
 package m2m
 
 import (
+	"io"
+	"net/http"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestM2MScopes(t *testing.T) {
@@ -38,3 +42,48 @@ func TestM2MScopes(t *testing.T) {
 		assert.Equal(t, []string{"all-apis", "my-scope"}, auth.scopes)
 	})
 }
+
+// recordingTransport answers every request with a fixed token body and keeps
+// the User-Agent it was sent with.
+type recordingTransport struct {
+	userAgents []string
+}
+
+func (rt *recordingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	rt.userAgents = append(rt.userAgents, req.Header.Get("User-Agent"))
+	return &http.Response{
+		StatusCode: http.StatusOK,
+		Header:     http.Header{"Content-Type": {"application/json"}},
+		Body:       io.NopCloser(strings.NewReader(`{"access_token":"tok","token_type":"Bearer","expires_in":3600}`)),
+		Request:    req,
+	}, nil
+}
+
+func TestAuthenticateUsesInjectedHTTPClient(t *testing.T) {
+	// Azure hosts resolve their endpoints offline, so no discovery request is made.
+	const host = "adb-1.2.azuredatabricks.net"
+
+	rt := &recordingTransport{}
+	auth := NewAuthenticator("id", "secret", host).(*authClient)
+	auth.SetHTTPClient(&http.Client{Transport: &userAgentTransport{base: rt, userAgent: "driver/1.0 (isv)"}})
+
+	req, _ := http.NewRequest(http.MethodPost, "https://"+host+"/sql/1.0/warehouses/x", nil)
+	require.NoError(t, auth.Authenticate(req))
+	assert.Equal(t, "Bearer tok", req.Header.Get("Authorization"))
+	assert.Equal(t, []string{"driver/1.0 (isv)"}, rt.userAgents)
+
+	// The cached token source is reused: no second token request.
+	req2, _ := http.NewRequest(http.MethodPost, "https://"+host+"/sql/1.0/warehouses/x", nil)
+	require.NoError(t, auth.Authenticate(req2))
+	assert.Len(t, rt.userAgents, 1)
+}
+
+type userAgentTransport struct {
+	base      http.RoundTripper
+	userAgent string
+}
+
+func (t *userAgentTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	req.Header.Set("User-Agent", t.userAgent)
+	return t.base.RoundTrip(req)
+}

auth/oauth/oauth.go

@@ -55,11 +55,14 @@ func GetEndpoint(ctx context.Context, hostName string) (oauth2.Endpoint, error)
 	// For normal workspace hosts this resolves to https://<host>/oidc, identical to
 	// the previous behavior.
 	//
-	// NOTE: this client uses the default transport, matching the existing
-	// oidc.NewProvider discovery below. A connector-supplied transport / TLS config
-	// (WithTransport, WithSkipTLSHostVerify) is not yet threaded into the OAuth
-	// endpoint-resolution path; that is a pre-existing limitation, tracked separately.
+	// The discovery client reuses the transport of an oauth2.HTTPClient supplied on
+	// ctx (the driver sets one carrying its User-Agent and WithTransport) so the
+	// lookup is attributable like the token requests that follow; without it the
+	// default transport is used.
 	client := &http.Client{Timeout: hostConfigTimeout}
+	if hc, ok := ctx.Value(oauth2.HTTPClient).(*http.Client); ok && hc != nil {
+		client.Transport = hc.Transport
+	}
 	issuerURL := resolveOIDCIssuer(ctx, client, hostName)
 	ctx = oidc.ClientContext(ctx, client)
 	ctx = oidc.InsecureIssuerURLContext(ctx, issuerURL)

auth/oauth/oauth_test.go

@@ -2,11 +2,15 @@ package oauth
 
 import (
 	"context"
+	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"reflect"
 	"strings"
 	"testing"
+
+	"golang.org/x/oauth2"
 )
 
 func TestInferCloudFromHost(t *testing.T) {
@@ -279,3 +283,47 @@ func TestFetchHostMetadata_failuresFallBack(t *testing.T) {
 		}
 	})
 }
+
+// TestGetEndpointUsesContextHTTPClientTransport asserts that discovery goes
+// through the transport of an oauth2.HTTPClient supplied on ctx: the fake
+// transport answers both the databricks-config lookup and the OIDC discovery
+// document, so no network is touched and the resolved endpoint is the one it served.
+func TestGetEndpointUsesContextHTTPClientTransport(t *testing.T) {
+	const host = "ws.cloud.databricks.com"
+	rt := &scriptedTransport{userAgents: map[string]string{}}
+	ctx := context.WithValue(context.Background(), oauth2.HTTPClient, &http.Client{Transport: rt})
+
+	endpoint, err := GetEndpoint(ctx, host)
+	if err != nil {
+		t.Fatalf("GetEndpoint: %v", err)
+	}
+	if want := "https://" + host + "/oidc/v1/token"; endpoint.TokenURL != want {
+		t.Fatalf("TokenURL = %q, want %q", endpoint.TokenURL, want)
+	}
+	for _, path := range []string{"/.well-known/databricks-config", "/oidc/.well-known/openid-configuration"} {
+		if _, ok := rt.userAgents[path]; !ok {
+			t.Fatalf("%s was not requested through the ctx transport (seen: %v)", path, rt.userAgents)
+		}
+	}
+}
+
+type scriptedTransport struct {
+	userAgents map[string]string
+}
+
+func (rt *scriptedTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	rt.userAgents[req.URL.Path] = req.Header.Get("User-Agent")
+	status, body := http.StatusNotFound, ""
+	if req.URL.Path == "/oidc/.well-known/openid-configuration" {
+		base := "https://" + req.URL.Host + "/oidc"
+		status = http.StatusOK
+		body = fmt.Sprintf(`{"issuer":%q,"authorization_endpoint":%q,"token_endpoint":%q,"jwks_uri":%q}`,
+			base, base+"/v1/authorize", base+"/v1/token", base+"/jwks.json")
+	}
+	return &http.Response{
+		StatusCode: status,
+		Header:     http.Header{"Content-Type": {"application/json"}},
+		Body:       io.NopCloser(strings.NewReader(body)),
+		Request:    req,
+	}, nil
+}

internal/client/client.go

@@ -622,25 +622,39 @@ func PooledClient(cfg *config.Config) *http.Client {
 		return nil
 	}
 
-	var tr *Transport
-	if cfg.Transport != nil {
-		tr = &Transport{
-			Base:  cfg.Transport,
-			Authr: cfg.Authenticator,
-		}
-	} else {
-		tr = &Transport{
-			Base:  PooledTransport(cfg),
-			Authr: cfg.Authenticator,
-		}
+	base := cfg.Transport
+	if base == nil {
+		base = PooledTransport(cfg)
+	}
+
+	// Authenticators that fetch their own tokens (OAuth M2M) get a client on the
+	// same base transport, with the driver User-Agent, so token traffic is
+	// attributed like Thrift traffic. It must not go through the auth Transport
+	// below, which would recurse into Authenticate.
+	if s, ok := cfg.Authenticator.(interface{ SetHTTPClient(*http.Client) }); ok {
+		s.SetHTTPClient(&http.Client{
+			Transport: &userAgentTransport{base: base, userAgent: BuildUserAgent(cfg)},
+			Timeout:   cfg.ClientTimeout,
+		})
 	}
 
 	return &http.Client{
-		Transport: tr,
+		Transport: &Transport{Base: base, Authr: cfg.Authenticator},
 		Timeout:   cfg.ClientTimeout,
 	}
 }
 
+type userAgentTransport struct {
+	base      http.RoundTripper
+	userAgent string
+}
+
+func (t *userAgentTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	req2 := cloneRequest(req)
+	req2.Header.Set("User-Agent", t.userAgent)
+	return t.base.RoundTrip(req2)
+}
+
 // cloneRequest returns a clone of the provided *http.Request.
 // The clone is a shallow copy of the struct and its Header map.
 func cloneRequest(r *http.Request) *http.Request {

internal/client/useragent_test.go

@@ -1,9 +1,12 @@
 package client
 
 import (
+	"net/http"
 	"testing"
 
+	"github.com/databricks/databricks-sql-go/auth/noop"
 	"github.com/databricks/databricks-sql-go/internal/config"
+	"github.com/stretchr/testify/require"
 )
 
 // clearAgentEnv blanks every env var agent.Detect inspects so the test
@@ -80,3 +83,49 @@ func TestBuildUserAgent(t *testing.T) {
 		}
 	})
 }
+
+type recordingAuth struct {
+	client *http.Client
+}
+
+func (a *recordingAuth) Authenticate(*http.Request) error { return nil }
+func (a *recordingAuth) SetHTTPClient(c *http.Client)     { a.client = c }
+
+type recordingRoundTripper struct {
+	userAgent string
+}
+
+func (rt *recordingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	rt.userAgent = req.Header.Get("User-Agent")
+	return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Request: req}, nil
+}
+
+func TestPooledClientInjectsUserAgentClientIntoAuthenticator(t *testing.T) {
+	clearAgentEnv(t)
+	base := &recordingRoundTripper{}
+	authr := &recordingAuth{}
+	cfg := &config.Config{
+		DriverName:    "godatabrickssqlconnector",
+		DriverVersion: "9.9.9",
+		UserConfig: config.UserConfig{
+			UserAgentEntry: "isv/product",
+			Authenticator:  authr,
+			Transport:      base,
+		},
+	}
+
+	PooledClient(cfg)
+
+	require.NotNil(t, authr.client, "authenticator should receive an http client")
+	req, _ := http.NewRequest(http.MethodPost, "https://host/oidc/token", nil)
+	resp, err := authr.client.Transport.RoundTrip(req)
+	require.NoError(t, err)
+	defer resp.Body.Close() //nolint:errcheck
+	require.Equal(t, BuildUserAgent(cfg), base.userAgent)
+	require.Empty(t, req.Header.Get("User-Agent"), "caller's request must not be mutated")
+}
+
+func TestPooledClientSkipsAuthenticatorsWithoutSetHTTPClient(t *testing.T) {
+	cfg := &config.Config{UserConfig: config.UserConfig{Authenticator: &noop.NoopAuth{}}}
+	require.NotNil(t, PooledClient(cfg))
+}