Add per-site authorize URL and per-domain token URL OAuth2 plumbing

Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>

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 {

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.<domain>).
+	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,
 		},
 	}

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

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

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

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

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