Add BuildTokenURLForSite for subdomain OAuth flows

Subdomain-based OAuth providers (e.g. Zendesk) authenticate against a
per-customer host, <subdomain>.zendesk.com, for both the authorize and
the token endpoint. The subdomain is known at initiate but, unlike
Datadog's region, the provider does not echo it back on the callback,
so it must survive the round-trip another way.

Carry the chosen site on the HMAC-signed OAuth state (OAuth2State.Site)
and add a BuildTokenURLForSite closure, symmetric with the existing
BuildAuthURLForSite, that builds the token URL from it at callback time.
The state is signed, so a tampered site is rejected before use; the
closure still re-validates the format. Datadog's domain-echo path is
unchanged and a provider sets at most one of the two closures.

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

pkg/connector/oauth2.go

@@ -75,13 +75,18 @@ 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 / BuildTokenURLForDomain / BuildTokenURLForSite
+		// are copied from the provider Registration by ApplyOAuth2Defaults.
+		// When set, the authorize URL is built per-site at initiate.
+		// BuildTokenURLForDomain builds the token URL from a host the provider
+		// echoes back on the callback (Datadog's `domain`);
+		// BuildTokenURLForSite builds it from the site carried in the signed
+		// state (Zendesk's subdomain), for providers that do not echo the host
+		// back. A provider sets at most one of the two. Nil for single-site
+		// providers.
 		BuildAuthURLForSite    func(site string) (string, error)
 		BuildTokenURLForDomain func(domain string) (string, error)
+		BuildTokenURLForSite   func(site string) (string, error)
 
 		// HTTPClient is used for the OAuth2 token-exchange request
 		// issued from CompleteWithState. It must be set by callers;
@@ -97,6 +102,16 @@ type (
 		ContinueURL     string   `json:"continue,omitempty"`
 		ConnectorID     string   `json:"cid,omitempty"` // Set when reconnecting an existing connector
 		RequestedScopes []string `json:"scopes,omitempty"`
+		// Site carries the per-customer site/subdomain chosen at initiate
+		// (opts.Site) to the callback for multi-site providers whose token
+		// host is NOT echoed back by the provider (e.g. Zendesk, whose
+		// token endpoint lives at <subdomain>.zendesk.com). It is signed
+		// into the state token, so a tampered value is rejected by the HMAC
+		// check; the callback still re-validates the format before using it
+		// to build a URL host. Empty for single-site providers and for
+		// multi-site providers that recover the host from a callback param
+		// (e.g. Datadog's `domain`).
+		Site string `json:"site,omitempty"`
 		// PKCENonce carries a random per-flow nonce between Initiate and
 		// Complete for providers that require PKCE. The actual
 		// code_verifier is DERIVED server-side from the state salt and this
@@ -194,6 +209,12 @@ func (c *OAuth2Connector) InitiateWithState(
 		return "", fmt.Errorf("cannot create state token: connector has no state signing key or client secret")
 	}
 
+	// Carry the per-customer site/subdomain (if any) into the signed state so
+	// it survives the round-trip to the callback. Multi-site providers whose
+	// token host the provider does not echo back (e.g. Zendesk) read it from
+	// the state at CompleteWithState. No-op (omitempty) when unset.
+	stateData.Site = opts.Site
+
 	// For PKCE providers a per-flow nonce is generated and stored in the
 	// signed state; the code verifier itself is DERIVED from salt+nonce
 	// (derivePKCEVerifier) and never serialized, so it stays secret even
@@ -371,7 +392,8 @@ func (c *OAuth2Connector) CompleteWithState(ctx context.Context, r *http.Request
 	}
 
 	tokenURL := c.TokenURL
-	if c.BuildTokenURLForDomain != nil {
+	switch {
+	case c.BuildTokenURLForDomain != nil:
 		domain := r.URL.Query().Get("domain")
 		if domain == "" {
 			return nil, nil, fmt.Errorf("cannot complete oauth2 flow: missing domain parameter")
@@ -382,6 +404,21 @@ func (c *OAuth2Connector) CompleteWithState(ctx context.Context, r *http.Request
 			return nil, nil, fmt.Errorf("cannot build token URL: %w", err)
 		}
 
+		tokenURL = built
+	case c.BuildTokenURLForSite != nil:
+		// The site/subdomain was carried in the signed state from initiate
+		// (the provider does not echo it back on the callback). The HMAC
+		// signature already authenticated the value; the closure re-validates
+		// the format before using it as a URL host.
+		if payload.Data.Site == "" {
+			return nil, nil, fmt.Errorf("cannot complete oauth2 flow: missing site in state")
+		}
+
+		built, err := c.BuildTokenURLForSite(payload.Data.Site)
+		if err != nil {
+			return nil, nil, fmt.Errorf("cannot build token URL: %w", err)
+		}
+
 		tokenURL = built
 	}
 
@@ -439,8 +476,8 @@ func (c *OAuth2Connector) CompleteWithState(ctx context.Context, r *http.Request
 	}
 
 	// Persist the per-customer token URL for multi-site providers so
-	// token refresh targets the same regional host (api.<domain>).
-	if c.BuildTokenURLForDomain != nil {
+	// token refresh targets the same regional/subdomain host.
+	if c.BuildTokenURLForDomain != nil || c.BuildTokenURLForSite != nil {
 		oauth2Conn.TokenURL = tokenURL
 	}
 

pkg/connector/provider/apply.go

@@ -49,6 +49,7 @@ func (r *Registry) ApplyOAuth2Defaults(p string, redirectURI string, c *connecto
 	c.RequiresPKCE = reg.RequiresPKCE
 	c.BuildAuthURLForSite = reg.BuildAuthURLForSite
 	c.BuildTokenURLForDomain = reg.BuildTokenURLForDomain
+	c.BuildTokenURLForSite = reg.BuildTokenURLForSite
 
 	// Deep copy ExtraAuthParams so per-connector mutations (e.g.
 	// incremental auth, scope overrides) cannot alias back into the

pkg/connector/provider/types.go

@@ -68,6 +68,13 @@ type Registration struct {
 	// 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)
+	// BuildTokenURLForSite builds the token endpoint URL from the
+	// per-customer site/subdomain carried in the signed OAuth state, for
+	// multi-site providers whose token host the provider does NOT echo back
+	// on the callback (e.g. Zendesk's <subdomain>.zendesk.com). It MUST
+	// validate site. A provider sets at most one of BuildTokenURLForDomain /
+	// BuildTokenURLForSite. Nil otherwise.
+	BuildTokenURLForSite func(site string) (string, error)
 
 	// Protocol support / GraphQL surface.
 	SupportsAPIKey            bool