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/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