From 082772465db38c22a1eadb18297467e7f120da51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aur=C3=A9lien=20Sibiril?= <81782+aureliensibiril@users.noreply.github.com> Date: Fri, 29 May 2026 14:48:30 +0200 Subject: [PATCH] Add per-site authorize URL and per-domain token URL OAuth2 plumbing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com> --- pkg/connector/connector.go | 5 + pkg/connector/oauth2.go | 69 ++++++- pkg/connector/oauth2_test.go | 186 ++++++++++++++++++ pkg/connector/provider/apply.go | 2 + pkg/connector/provider/apply_test.go | 33 ++++ pkg/connector/provider/types.go | 9 + .../api/console/v1/connector_initiate.go | 2 +- 7 files changed, 296 insertions(+), 10 deletions(-) diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go index 34a425d38..c60605e27 100644 --- a/pkg/connector/connector.go +++ b/pkg/connector/connector.go @@ -39,6 +39,11 @@ type ( // existing connector: the callback updates the row in place // instead of creating a new one. ConnectorID string + // Site selects a per-customer region/site for multi-site + // providers (e.g. Datadog). Consumed by the connector's + // Registration.BuildAuthURLForSite. Empty for single-site + // providers. + Site string } Connector interface { diff --git a/pkg/connector/oauth2.go b/pkg/connector/oauth2.go index 2101c4c57..7043f8540 100644 --- a/pkg/connector/oauth2.go +++ b/pkg/connector/oauth2.go @@ -75,6 +75,14 @@ type ( // wiring, never serialized. StateSigningKey string + // BuildAuthURLForSite / BuildTokenURLForDomain are copied from + // the provider Registration by ApplyOAuth2Defaults. When set, + // the authorize URL is built per-site at initiate and the token + // URL per-domain at callback (multi-site providers, e.g. + // Datadog). Nil for single-site providers. + BuildAuthURLForSite func(site string) (string, error) + BuildTokenURLForDomain func(domain string) (string, error) + // HTTPClient is used for the OAuth2 token-exchange request // issued from CompleteWithState. It must be set by callers; // (*provider.Registry).ApplyOAuth2Defaults assigns an @@ -236,7 +244,21 @@ func (c *OAuth2Connector) InitiateWithState( authCodeQuery.Set(k, v) } - u, err := url.Parse(c.AuthURL) + authURL := c.AuthURL + if c.BuildAuthURLForSite != nil { + if opts.Site == "" { + return "", fmt.Errorf("cannot initiate connector: site is required for multi-site providers") + } + + built, err := c.BuildAuthURLForSite(opts.Site) + if err != nil { + return "", fmt.Errorf("cannot build auth URL for site: %w", err) + } + + authURL = built + } + + u, err := url.Parse(authURL) if err != nil { return "", fmt.Errorf("cannot parse auth URL: %w", err) } @@ -348,7 +370,22 @@ func (c *OAuth2Connector) CompleteWithState(ctx context.Context, r *http.Request codeVerifier = derivePKCEVerifier(c.stateSalt(), payload.Data.PKCENonce) } - tokenRequest, err := c.buildTokenRequest(ctx, code, c.RedirectURI, codeVerifier) + tokenURL := c.TokenURL + if c.BuildTokenURLForDomain != nil { + domain := r.URL.Query().Get("domain") + if domain == "" { + return nil, nil, fmt.Errorf("cannot complete oauth2 flow: missing domain parameter") + } + + built, err := c.BuildTokenURLForDomain(domain) + if err != nil { + return nil, nil, fmt.Errorf("cannot build token URL: %w", err) + } + + tokenURL = built + } + + tokenRequest, err := c.buildTokenRequest(ctx, code, c.RedirectURI, codeVerifier, tokenURL) if err != nil { return nil, nil, err } @@ -401,6 +438,12 @@ func (c *OAuth2Connector) CompleteWithState(ctx context.Context, r *http.Request oauth2Conn.ExpiresAt = time.Now().Add(time.Duration(rawToken.ExpiresIn) * time.Second) } + // Persist the per-customer token URL for multi-site providers so + // token refresh targets the same regional host (api.). + if c.BuildTokenURLForDomain != nil { + oauth2Conn.TokenURL = tokenURL + } + if payload.Data.Provider == SlackProvider { conn, _, err := ParseSlackTokenResponse(body, oauth2Conn, organizationID) return conn, &payload.Data, err @@ -436,11 +479,11 @@ func basicAuthHeader(clientID, clientSecret string) string { // endpoint with the headers shared by the form-body auth methods // ("basic-form", "post-form", and the public-client "none"). Callers set any // extra auth header (e.g. Basic) on the returned request. -func (c *OAuth2Connector) newFormTokenRequest(ctx context.Context, form url.Values) (*http.Request, error) { +func (c *OAuth2Connector) newFormTokenRequest(ctx context.Context, form url.Values, tokenURL string) (*http.Request, error) { req, err := http.NewRequestWithContext( ctx, http.MethodPost, - c.TokenURL, + tokenURL, strings.NewReader(form.Encode()), ) if err != nil { @@ -458,7 +501,7 @@ func (c *OAuth2Connector) newFormTokenRequest(ctx context.Context, form url.Valu // on c.TokenEndpointAuth to support different provider requirements. When // codeVerifier is non-empty (PKCE-enabled providers), it is replayed as // `code_verifier` in the request body. -func (c *OAuth2Connector) buildTokenRequest(ctx context.Context, code, redirectURI, codeVerifier string) (*http.Request, error) { +func (c *OAuth2Connector) buildTokenRequest(ctx context.Context, code, redirectURI, codeVerifier, tokenURL string) (*http.Request, error) { switch c.TokenEndpointAuth { case "basic-json": // JSON body with Basic auth header (Notion). @@ -479,7 +522,7 @@ func (c *OAuth2Connector) buildTokenRequest(ctx context.Context, code, redirectU req, err := http.NewRequestWithContext( ctx, http.MethodPost, - c.TokenURL, + tokenURL, bytes.NewReader(jsonBody), ) if err != nil { @@ -504,7 +547,7 @@ func (c *OAuth2Connector) buildTokenRequest(ctx context.Context, code, redirectU formData.Set("code_verifier", codeVerifier) } - req, err := c.newFormTokenRequest(ctx, formData) + req, err := c.newFormTokenRequest(ctx, formData, tokenURL) if err != nil { return nil, err } @@ -533,7 +576,7 @@ func (c *OAuth2Connector) buildTokenRequest(ctx context.Context, code, redirectU formData.Set("code_verifier", codeVerifier) } - return c.newFormTokenRequest(ctx, formData) + return c.newFormTokenRequest(ctx, formData, tokenURL) } } @@ -598,11 +641,19 @@ func (c *OAuth2Connection) RefreshableClient(ctx context.Context, cfg OAuth2Refr authStyle = oauth2.AuthStyleInHeader } + // Multi-site providers persist a per-customer token URL on the + // connection; prefer it over the static registration TokenURL so + // refresh targets the correct regional host. + tokenURL := cfg.TokenURL + if c.TokenURL != "" { + tokenURL = c.TokenURL + } + config := &oauth2.Config{ ClientID: cfg.ClientID, ClientSecret: cfg.ClientSecret, Endpoint: oauth2.Endpoint{ - TokenURL: cfg.TokenURL, + TokenURL: tokenURL, AuthStyle: authStyle, }, } diff --git a/pkg/connector/oauth2_test.go b/pkg/connector/oauth2_test.go index bf270bf28..a5501ef4d 100644 --- a/pkg/connector/oauth2_test.go +++ b/pkg/connector/oauth2_test.go @@ -18,6 +18,7 @@ import ( "context" "encoding/base64" "encoding/json" + "fmt" "io" "net/http" "net/http/httptest" @@ -50,6 +51,7 @@ func TestBuildTokenRequest_PostForm(t *testing.T) { "test-code", "https://example.com/callback", "", + connector.TokenURL, ) require.NoError(t, err) @@ -86,6 +88,7 @@ func TestBuildTokenRequest_PostForm(t *testing.T) { "test-code", "https://example.com/callback", "", + connector.TokenURL, ) require.NoError(t, err) @@ -121,6 +124,7 @@ func TestBuildTokenRequest_BasicForm(t *testing.T) { "test-code", "https://example.com/callback", "", + connector.TokenURL, ) require.NoError(t, err) @@ -164,6 +168,7 @@ func TestBuildTokenRequest_BasicJSON(t *testing.T) { "test-code", "https://example.com/callback", "", + connector.TokenURL, ) require.NoError(t, err) @@ -704,6 +709,187 @@ func TestInitiateWithState_PKCE(t *testing.T) { }) } +func TestInitiateWithState_PerSiteAuthURL(t *testing.T) { + t.Parallel() + + c := &OAuth2Connector{ + ClientID: "cid", + ClientSecret: "secret", + RedirectURI: "https://probo.example/cb", + RequiresPKCE: true, + BuildAuthURLForSite: DatadogAuthorizeURL, + } + + got, err := c.InitiateWithState(context.Background(), + OAuth2State{OrganizationID: "org", Provider: DatadogProvider}, + InitiateOptions{Scopes: []string{"user_access_read"}, Site: "US3"}, + ) + require.NoError(t, err) + + u, err := url.Parse(got) + require.NoError(t, err) + assert.Equal(t, "us3.datadoghq.com", u.Host) + assert.Equal(t, "/oauth2/v1/authorize", u.Path) + assert.NotEmpty(t, u.Query().Get("code_challenge")) +} + +func TestInitiateWithState_MissingSiteForMultiSite(t *testing.T) { + t.Parallel() + + c := &OAuth2Connector{ + ClientID: "cid", + ClientSecret: "secret", + BuildAuthURLForSite: DatadogAuthorizeURL, + } + + _, err := c.InitiateWithState(context.Background(), + OAuth2State{OrganizationID: "org", Provider: DatadogProvider}, + InitiateOptions{Site: ""}, + ) + require.Error(t, err) +} + +func TestInitiateWithState_InvalidSiteRejected(t *testing.T) { + t.Parallel() + + c := &OAuth2Connector{ + ClientID: "cid", + ClientSecret: "secret", + BuildAuthURLForSite: DatadogAuthorizeURL, + } + + _, err := c.InitiateWithState(context.Background(), + OAuth2State{OrganizationID: "org", Provider: DatadogProvider}, + InitiateOptions{Site: "BOGUS"}, + ) + require.Error(t, err) +} + +func TestCompleteWithState_PerDomainTokenURL(t *testing.T) { + t.Parallel() + + var gotPath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"access_token":"at","refresh_token":"rt","expires_in":3600,"token_type":"Bearer"}`)) + })) + defer srv.Close() + + // Build a token-URL closure that targets the httptest server, + // mirroring DatadogTokenURL's shape (validate then build). + c := &OAuth2Connector{ + ClientID: "cid", + ClientSecret: "secret", + RedirectURI: "https://probo.example/cb", + RequiresPKCE: true, + HTTPClient: httpclient.DefaultClient(httpclient.WithSSRFProtection(), httpclient.WithSSRFAllowLoopback()), + BuildTokenURLForDomain: func(domain string) (string, error) { + if domain != "us3.datadoghq.com" { + return "", fmt.Errorf("unknown domain") + } + return srv.URL + "/oauth2/v1/token", nil + }, + } + + state, err := statelesstoken.NewToken(c.ClientSecret, OAuth2TokenType, OAuth2TokenTTL, + OAuth2State{OrganizationID: validOrgGID(t), Provider: DatadogProvider}) + require.NoError(t, err) + + req := httptest.NewRequest(http.MethodGet, + "https://probo.example/cb?code=abc&state="+state+"&domain=us3.datadoghq.com", nil) + + conn, _, err := c.CompleteWithState(context.Background(), req) + require.NoError(t, err) + assert.Equal(t, "/oauth2/v1/token", gotPath) + + oc, ok := conn.(*OAuth2Connection) + require.True(t, ok) + assert.Equal(t, srv.URL+"/oauth2/v1/token", oc.TokenURL) +} + +func TestCompleteWithState_MissingDomainForMultiSite(t *testing.T) { + t.Parallel() + + c := &OAuth2Connector{ + ClientID: "cid", + ClientSecret: "secret", + RedirectURI: "https://probo.example/cb", + HTTPClient: httpclient.DefaultClient(httpclient.WithSSRFProtection(), httpclient.WithSSRFAllowLoopback()), + BuildTokenURLForDomain: func(string) (string, error) { return "", fmt.Errorf("unused") }, + } + + state, err := statelesstoken.NewToken(c.ClientSecret, OAuth2TokenType, OAuth2TokenTTL, + OAuth2State{OrganizationID: validOrgGID(t), Provider: DatadogProvider}) + require.NoError(t, err) + + req := httptest.NewRequest(http.MethodGet, + "https://probo.example/cb?code=abc&state="+state, nil) + + _, _, err = c.CompleteWithState(context.Background(), req) + require.Error(t, err) +} + +// TestCompleteWithState_InvalidDomainRejected exercises the SSRF guard: a +// tampered callback `domain` must fail the flow (the closure validates against +// the fixed allow-list) before credentials are POSTed anywhere. +func TestCompleteWithState_InvalidDomainRejected(t *testing.T) { + t.Parallel() + + c := &OAuth2Connector{ + ClientID: "cid", + ClientSecret: "secret", + RedirectURI: "https://probo.example/cb", + HTTPClient: httpclient.DefaultClient(httpclient.WithSSRFProtection(), httpclient.WithSSRFAllowLoopback()), + BuildTokenURLForDomain: DatadogTokenURL, + } + + state, err := statelesstoken.NewToken(c.ClientSecret, OAuth2TokenType, OAuth2TokenTTL, + OAuth2State{OrganizationID: validOrgGID(t), Provider: DatadogProvider}) + require.NoError(t, err) + + req := httptest.NewRequest(http.MethodGet, + "https://probo.example/cb?code=abc&state="+state+"&domain=evil.example.com", nil) + + _, _, err = c.CompleteWithState(context.Background(), req) + require.Error(t, err) +} + +func TestRefreshableClient_PrefersConnectionTokenURL(t *testing.T) { + t.Parallel() + + var gotHost string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotHost = r.Host + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"access_token":"new","token_type":"Bearer","expires_in":3600}`)) + })) + defer srv.Close() + + conn := &OAuth2Connection{ + AccessToken: "old", + RefreshToken: "rt", + TokenType: "Bearer", + ExpiresAt: time.Now().Add(-time.Hour), + TokenURL: srv.URL, // per-connection (Datadog-style) + } + + // cfg.TokenURL is empty (multi-site providers carry no static token URL). + _, err := conn.RefreshableClient(context.Background(), OAuth2RefreshConfig{ + ClientID: "cid", ClientSecret: "secret", + }, httpclient.WithSSRFAllowLoopback()) + require.NoError(t, err) + + u, _ := url.Parse(srv.URL) + assert.Equal(t, u.Host, gotHost) + assert.Equal(t, "new", conn.AccessToken) +} + +func validOrgGID(t *testing.T) string { + t.Helper() + return gid.New(gid.NewTenantID(), 0).String() +} + // TestGeneratePKCENonce exercises the nonce generator: each call must // return a fresh value, encoded as RFC 4648 §5 base64url-without-padding // (32 bytes yields 43 chars). The nonce seeds derivePKCEVerifier, so a diff --git a/pkg/connector/provider/apply.go b/pkg/connector/provider/apply.go index 85ac9e9bb..5c25b061c 100644 --- a/pkg/connector/provider/apply.go +++ b/pkg/connector/provider/apply.go @@ -47,6 +47,8 @@ func (r *Registry) ApplyOAuth2Defaults(p string, redirectURI string, c *connecto c.TokenEndpointAuth = reg.TokenEndpointAuth c.SupportsIncrementalAuth = reg.SupportsIncrementalAuth c.RequiresPKCE = reg.RequiresPKCE + c.BuildAuthURLForSite = reg.BuildAuthURLForSite + c.BuildTokenURLForDomain = reg.BuildTokenURLForDomain // Deep copy ExtraAuthParams so per-connector mutations (e.g. // incremental auth, scope overrides) cannot alias back into the diff --git a/pkg/connector/provider/apply_test.go b/pkg/connector/provider/apply_test.go index 56defdcc6..82d708c1e 100644 --- a/pkg/connector/provider/apply_test.go +++ b/pkg/connector/provider/apply_test.go @@ -15,13 +15,18 @@ package provider_test import ( + "context" + "net/http" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.gearno.de/kit/log" + "go.probo.inc/probo/pkg/accessreview/drivers" "go.probo.inc/probo/pkg/connector" "go.probo.inc/probo/pkg/connector/provider" + "go.probo.inc/probo/pkg/coredata" ) // TestApplyOAuth2Defaults_AuthURLFromSlug verifies that providers whose @@ -113,3 +118,31 @@ func TestApplyOAuth2Defaults_PublicClientTokenAuth(t *testing.T) { "PostHog must use token_endpoint_auth_method none (public client)") assert.True(t, c.RequiresPKCE, "PostHog public client must require PKCE") } + +// TestApplyOAuth2Defaults_CopiesSiteClosures verifies the multi-site +// per-provider closures (BuildAuthURLForSite, BuildTokenURLForDomain) are +// copied from the Registration onto the OAuth2Connector. +func TestApplyOAuth2Defaults_CopiesSiteClosures(t *testing.T) { + t.Parallel() + + r := provider.NewRegistry() + // Uses the PagerDuty enum (already exists) so Task 2 builds and commits + // independently of Task 3. The closures themselves are Datadog's, from + // Task 1 — this only asserts ApplyOAuth2Defaults copies them through. + require.NoError(t, r.Register(&provider.Registration{ + Provider: coredata.ConnectorProviderPagerDuty, + DisplayName: "PagerDuty", + OAuth2Scopes: []string{"users.read"}, + RequiresPKCE: true, + BuildAuthURLForSite: connector.DatadogAuthorizeURL, + BuildTokenURLForDomain: connector.DatadogTokenURL, + NewDriver: func(context.Context, *http.Client, *coredata.Connector, *log.Logger) (drivers.Driver, error) { + return nil, nil + }, + })) + + var c connector.OAuth2Connector + require.NoError(t, r.ApplyOAuth2Defaults("PAGERDUTY", "https://probo.example/cb", &c)) + require.NotNil(t, c.BuildAuthURLForSite) + require.NotNil(t, c.BuildTokenURLForDomain) +} diff --git a/pkg/connector/provider/types.go b/pkg/connector/provider/types.go index 44a65ad89..a556f3219 100644 --- a/pkg/connector/provider/types.go +++ b/pkg/connector/provider/types.go @@ -59,6 +59,15 @@ type Registration struct { // it as a path segment. It must construct the URL with net/url and // escape the slug. Nil for providers with a fully static AuthURL. BuildAuthURL func(slug string) (string, error) + // BuildAuthURLForSite builds the authorize URL for a per-customer + // site supplied at initiate time (multi-site providers, e.g. + // Datadog). It MUST validate site against a fixed allow-list and + // construct the URL with net/url. Nil for single-site providers. + BuildAuthURLForSite func(site string) (string, error) + // BuildTokenURLForDomain builds the token endpoint URL from the API + // domain the provider returns on the OAuth callback (multi-site + // providers, e.g. Datadog). It MUST validate domain. Nil otherwise. + BuildTokenURLForDomain func(domain string) (string, error) // Protocol support / GraphQL surface. SupportsAPIKey bool diff --git a/pkg/server/api/console/v1/connector_initiate.go b/pkg/server/api/console/v1/connector_initiate.go index d618f7483..ec752821d 100644 --- a/pkg/server/api/console/v1/connector_initiate.go +++ b/pkg/server/api/console/v1/connector_initiate.go @@ -112,7 +112,7 @@ func handleConnectorInitiate( // Union-not-delta because most providers replace rather than // merge. No short-circuit: every reconnect runs the full OAuth // flow so revoked or stale tokens are never silently reused. - opts := connector.InitiateOptions{Scopes: requestedScopes} + opts := connector.InitiateOptions{Scopes: requestedScopes, Site: r.URL.Query().Get("site")} if existing != nil { opts.Scopes = connector.UnionScopes(existing.Connection.Scopes(), requestedScopes) opts.IncludeGrantedScopes = true