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)) +}