From ef4bcd762176072196b7180bba5aee6e4f55fa2f Mon Sep 17 00:00:00 2001 From: Sevban Bayrak Date: Thu, 8 Oct 2026 15:04:11 +0300 Subject: [PATCH] fix(auth): thread connector transport and User-Agent into OAuth M2M token requests With WithClientCredentials the token exchange (and the OIDC discovery before it) was issued through http.DefaultClient, so system.access.audit recorded mintOAuthToken and workspaceInHouseOAuthClientAuthentication as "Go-http-client/2.0" instead of the driver User-Agent and WithUserAgentEntry. The M2M authenticator is built when its option is applied, before the User-Agent entry or transport may be set, so the driver now injects an *http.Client once the config is final: PooledClient passes a client on the same base transport, with BuildUserAgent, to any authenticator exposing SetHTTPClient. The M2M authenticator stores it on the oauth2 context used for GetConfig and TokenSource, and GetEndpoint reuses that client's transport for discovery. Authenticators without SetHTTPClient are unaffected. Fixes #413 Signed-off-by: Sevban Bayrak --- CHANGELOG.md | 1 + auth/oauth/m2m/m2m.go | 20 +++++++++++-- auth/oauth/m2m/m2m_test.go | 49 +++++++++++++++++++++++++++++++ auth/oauth/oauth.go | 11 ++++--- auth/oauth/oauth_test.go | 48 ++++++++++++++++++++++++++++++ internal/client/client.go | 38 ++++++++++++++++-------- internal/client/useragent_test.go | 49 +++++++++++++++++++++++++++++++ 7 files changed, 198 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a4e90d9c..27e024f4 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/auth/oauth/m2m/m2m.go b/auth/oauth/m2m/m2m.go index 76b5e513..ed5af607 100644 --- a/auth/oauth/m2m/m2m.go +++ b/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") diff --git a/auth/oauth/m2m/m2m_test.go b/auth/oauth/m2m/m2m_test.go index 1ff8ce97..1600bd82 100644 --- a/auth/oauth/m2m/m2m_test.go +++ b/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) +} diff --git a/auth/oauth/oauth.go b/auth/oauth/oauth.go index d0fd3f66..e43fc325 100644 --- a/auth/oauth/oauth.go +++ b/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:///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) diff --git a/auth/oauth/oauth_test.go b/auth/oauth/oauth_test.go index 96d71a2f..5fec9e11 100644 --- a/auth/oauth/oauth_test.go +++ b/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 +} diff --git a/internal/client/client.go b/internal/client/client.go index fec632cd..89226388 100644 --- a/internal/client/client.go +++ b/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 { diff --git a/internal/client/useragent_test.go b/internal/client/useragent_test.go index f662ce6f..07a3bf80 100644 --- a/internal/client/useragent_test.go +++ b/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)) +}