feat(connector): support OAuth2 scope preservation and incremental auth

Extend the OAuth2 connector so a reconnect can request the union of previously granted and newly requested scopes without losing either. Four related changes: - Connection gains Scopes() []string so callers no longer need a type switch to reach the scope set. OAuth2Connection and APIKeyConnection implement it; SlackConnection inherits via embedding. - OAuth2State carries RequestedScopes and CompleteWithState falls back to it when the provider omits the scope field (RFC 6749 §5.1 allows this when granted equals requested). Without the fallback the stored Scope would be empty and the next reconnect would have no diff base. - providerDefinition gains SupportsIncrementalAuth, set only for Google Workspace. When the flag is true and the caller passes InitiateOptions.IncludeGrantedScopes, the auth URL carries include_granted_scopes=true and the prompt=consent param is dropped so reuse flows see only the delta consent screen. - InitiateOptions gains ConnectorID so the reconnect case is passed explicitly instead of relying on the caller to mutate r.URL.Query. Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>
2026-04-09 01:02:49 +02:00
parent 62bab0f732
commit 71f6364df0
5 changed files with 304 additions and 22 deletions
--- a/pkg/connector/apikey.go
+++ b/pkg/connector/apikey.go
@@ -32,6 +32,10 @@ func (c *APIKeyConnection) Type() ProtocolType {
 	return ProtocolAPIKey
 }

+func (c *APIKeyConnection) Scopes() []string {
+	return nil
+}
+
 func (c *APIKeyConnection) Client(ctx context.Context) (*http.Client, error) {
 	transport := &oauth2Transport{
 		token:      c.APIKey,
--- a/pkg/connector/connector.go
+++ b/pkg/connector/connector.go
@@ -32,6 +32,13 @@ type (
 	// a different set of scopes (e.g. SCIM bridge vs access review).
 	InitiateOptions struct {
 		Scopes []string
+		// IncludeGrantedScopes is honored only when the provider has
+		// SupportsIncrementalAuth=true.
+		IncludeGrantedScopes bool
+		// ConnectorID, when set, marks this flow as a reconnect of an
+		// existing connector: the callback updates the row in place
+		// instead of creating a new one.
+		ConnectorID string
 	}

 	Connector interface {
@@ -42,6 +49,7 @@ type (
 	Connection interface {
 		Type() ProtocolType
 		Client(ctx context.Context) (*http.Client, error)
+		Scopes() []string

 		json.Unmarshaler
 		json.Marshaler
--- a/pkg/connector/oauth2.go
+++ b/pkg/connector/oauth2.go
@@ -41,20 +41,22 @@ import (

 type (
 	OAuth2Connector struct {
-		ClientID          string
-		ClientSecret      string
-		RedirectURI       string
-		AuthURL           string
-		TokenURL          string
-		ExtraAuthParams   map[string]string // Optional: extra params for auth URL (e.g., access_type=offline for Google)
-		TokenEndpointAuth string            // "post-form" (default), "basic-form", or "basic-json"
+		ClientID                string
+		ClientSecret            string
+		RedirectURI             string
+		AuthURL                 string
+		TokenURL                string
+		ExtraAuthParams         map[string]string // Optional: extra params for auth URL (e.g., access_type=offline for Google)
+		TokenEndpointAuth       string            // "post-form" (default), "basic-form", or "basic-json"
+		SupportsIncrementalAuth bool
 	}

 	OAuth2State struct {
-		OrganizationID string `json:"oid"`
-		Provider       string `json:"provider"`
-		ContinueURL    string `json:"continue,omitempty"`
-		ConnectorID    string `json:"cid,omitempty"` // Set when reconnecting an existing connector
+		OrganizationID  string   `json:"oid"`
+		Provider        string   `json:"provider"`
+		ContinueURL     string   `json:"continue,omitempty"`
+		ConnectorID     string   `json:"cid,omitempty"` // Set when reconnecting an existing connector
+		RequestedScopes []string `json:"scopes,omitempty"`
 	}

 	OAuth2Connection struct {
@@ -105,16 +107,15 @@ func (c *OAuth2Connector) Initiate(
 	r *http.Request,
 ) (string, error) {
 	stateData := OAuth2State{
-		OrganizationID: organizationID.String(),
-		Provider:       provider,
+		OrganizationID:  organizationID.String(),
+		Provider:        provider,
+		ConnectorID:     opts.ConnectorID,
+		RequestedScopes: opts.Scopes,
 	}
 	if r != nil {
 		if continueURL := r.URL.Query().Get("continue"); continueURL != "" {
 			stateData.ContinueURL = continueURL
 		}
-		if connectorID := r.URL.Query().Get("connector_id"); connectorID != "" {
-			stateData.ConnectorID = connectorID
-		}
 	}
 	return c.InitiateWithState(ctx, stateData, opts, r)
 }
@@ -141,8 +142,18 @@ func (c *OAuth2Connector) InitiateWithState(
 		authCodeQuery.Set("scope", strings.Join(opts.Scopes, " "))
 	}

-	// Add any extra auth params (e.g., access_type=offline, prompt=consent for Google)
+	incrementalAuth := c.SupportsIncrementalAuth && opts.IncludeGrantedScopes
+	if incrementalAuth {
+		authCodeQuery.Set("include_granted_scopes", "true")
+	}
+
+	// Skip prompt=consent when doing incremental auth so the user sees
+	// only the delta, not a full re-consent. First-install flows keep it
+	// because IncludeGrantedScopes is false there.
 	for k, v := range c.ExtraAuthParams {
+		if incrementalAuth && k == "prompt" && v == "consent" {
+			continue
+		}
 		authCodeQuery.Set(k, v)
 	}

@@ -225,11 +236,19 @@ func (c *OAuth2Connector) CompleteWithState(ctx context.Context, r *http.Request
 		return nil, nil, fmt.Errorf("cannot decode token response: %w", err)
 	}

+	grantedScope := rawToken.Scope
+	if grantedScope == "" {
+		// RFC 6749 §5.1: scope is OPTIONAL when identical to the
+		// requested scope. Fall back to what we asked for so
+		// subsequent reconnect diffs have a meaningful base.
+		grantedScope = FormatScopeString(payload.Data.RequestedScopes)
+	}
+
 	oauth2Conn := OAuth2Connection{
 		AccessToken:  rawToken.AccessToken,
 		RefreshToken: rawToken.RefreshToken,
 		TokenType:    rawToken.TokenType,
-		Scope:        rawToken.Scope,
+		Scope:        grantedScope,
 	}

 	// Convert expires_in (seconds) to expires_at (absolute time)
@@ -335,6 +354,10 @@ func (c *OAuth2Connection) Type() ProtocolType {
 	return ProtocolOAuth2
 }

+func (c *OAuth2Connection) Scopes() []string {
+	return ParseScopeString(c.Scope)
+}
+
 func (c *OAuth2Connection) Client(ctx context.Context) (*http.Client, error) {
 	return c.ClientWithOptions(ctx)
 }
--- a/pkg/connector/oauth2_test.go
+++ b/pkg/connector/oauth2_test.go
@@ -28,6 +28,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.probo.inc/probo/pkg/gid"
+	"go.probo.inc/probo/pkg/statelesstoken"
 )

 func TestBuildTokenRequest_PostForm(t *testing.T) {
@@ -310,4 +311,247 @@ func TestInitiateWithState_Scopes(t *testing.T) {
 		require.NoError(t, err)
 		assert.False(t, parsed.Query().Has("scope"), "scope param should be absent when no scopes provided")
 	})
+
+	t.Run("include_granted_scopes set when provider supports and caller requests", func(t *testing.T) {
+		t.Parallel()
+
+		c := &OAuth2Connector{
+			ClientID:                "id",
+			ClientSecret:            "secret",
+			RedirectURI:             "https://example.com/cb",
+			AuthURL:                 "https://provider.example.com/authorize",
+			SupportsIncrementalAuth: true,
+		}
+
+		orgID := gid.New(gid.NewTenantID(), 0)
+
+		u, err := c.InitiateWithState(
+			context.Background(),
+			OAuth2State{OrganizationID: orgID.String(), Provider: "TEST"},
+			InitiateOptions{
+				Scopes:               []string{"read:user"},
+				IncludeGrantedScopes: true,
+			},
+			nil,
+		)
+		require.NoError(t, err)
+
+		parsed, err := url.Parse(u)
+		require.NoError(t, err)
+		assert.Equal(t, "true", parsed.Query().Get("include_granted_scopes"))
+	})
+
+	t.Run("include_granted_scopes absent when provider does not support it", func(t *testing.T) {
+		t.Parallel()
+
+		c := &OAuth2Connector{
+			ClientID:                "id",
+			ClientSecret:            "secret",
+			RedirectURI:             "https://example.com/cb",
+			AuthURL:                 "https://provider.example.com/authorize",
+			SupportsIncrementalAuth: false,
+		}
+
+		orgID := gid.New(gid.NewTenantID(), 0)
+
+		u, err := c.InitiateWithState(
+			context.Background(),
+			OAuth2State{OrganizationID: orgID.String(), Provider: "TEST"},
+			InitiateOptions{
+				Scopes:               []string{"read:user"},
+				IncludeGrantedScopes: true,
+			},
+			nil,
+		)
+		require.NoError(t, err)
+
+		parsed, err := url.Parse(u)
+		require.NoError(t, err)
+		assert.False(t, parsed.Query().Has("include_granted_scopes"))
+	})
+
+	t.Run("include_granted_scopes absent when caller does not request", func(t *testing.T) {
+		t.Parallel()
+
+		c := &OAuth2Connector{
+			ClientID:                "id",
+			ClientSecret:            "secret",
+			RedirectURI:             "https://example.com/cb",
+			AuthURL:                 "https://provider.example.com/authorize",
+			SupportsIncrementalAuth: true,
+		}
+
+		orgID := gid.New(gid.NewTenantID(), 0)
+
+		u, err := c.InitiateWithState(
+			context.Background(),
+			OAuth2State{OrganizationID: orgID.String(), Provider: "TEST"},
+			InitiateOptions{Scopes: []string{"read:user"}},
+			nil,
+		)
+		require.NoError(t, err)
+
+		parsed, err := url.Parse(u)
+		require.NoError(t, err)
+		assert.False(t, parsed.Query().Has("include_granted_scopes"))
+	})
+
+	t.Run("prompt=consent skipped when incremental auth is active", func(t *testing.T) {
+		t.Parallel()
+
+		c := &OAuth2Connector{
+			ClientID:                "id",
+			ClientSecret:            "secret",
+			RedirectURI:             "https://example.com/cb",
+			AuthURL:                 "https://provider.example.com/authorize",
+			SupportsIncrementalAuth: true,
+			ExtraAuthParams: map[string]string{
+				"access_type": "offline",
+				"prompt":      "consent",
+			},
+		}
+
+		orgID := gid.New(gid.NewTenantID(), 0)
+
+		u, err := c.InitiateWithState(
+			context.Background(),
+			OAuth2State{OrganizationID: orgID.String(), Provider: "TEST"},
+			InitiateOptions{
+				Scopes:               []string{"read:user"},
+				IncludeGrantedScopes: true,
+			},
+			nil,
+		)
+		require.NoError(t, err)
+
+		parsed, err := url.Parse(u)
+		require.NoError(t, err)
+		assert.Equal(t, "offline", parsed.Query().Get("access_type"))
+		assert.False(t, parsed.Query().Has("prompt"), "prompt=consent should be skipped when doing incremental auth on a provider that supports it")
+		assert.Equal(t, "true", parsed.Query().Get("include_granted_scopes"))
+	})
+
+	t.Run("prompt=consent preserved on first install", func(t *testing.T) {
+		t.Parallel()
+
+		c := &OAuth2Connector{
+			ClientID:                "id",
+			ClientSecret:            "secret",
+			RedirectURI:             "https://example.com/cb",
+			AuthURL:                 "https://provider.example.com/authorize",
+			SupportsIncrementalAuth: true,
+			ExtraAuthParams: map[string]string{
+				"access_type": "offline",
+				"prompt":      "consent",
+			},
+		}
+
+		orgID := gid.New(gid.NewTenantID(), 0)
+
+		u, err := c.InitiateWithState(
+			context.Background(),
+			OAuth2State{OrganizationID: orgID.String(), Provider: "TEST"},
+			InitiateOptions{
+				Scopes:               []string{"read:user"},
+				IncludeGrantedScopes: false, // first install, no existing grant
+			},
+			nil,
+		)
+		require.NoError(t, err)
+
+		parsed, err := url.Parse(u)
+		require.NoError(t, err)
+		assert.Equal(t, "offline", parsed.Query().Get("access_type"))
+		assert.Equal(t, "consent", parsed.Query().Get("prompt"), "prompt=consent must still fire on first install so Google issues a refresh token")
+		assert.False(t, parsed.Query().Has("include_granted_scopes"))
+	})
+
+	t.Run("prompt=consent preserved when provider does not support incremental auth", func(t *testing.T) {
+		t.Parallel()
+
+		c := &OAuth2Connector{
+			ClientID:                "id",
+			ClientSecret:            "secret",
+			RedirectURI:             "https://example.com/cb",
+			AuthURL:                 "https://provider.example.com/authorize",
+			SupportsIncrementalAuth: false,
+			ExtraAuthParams: map[string]string{
+				"prompt": "consent",
+			},
+		}
+
+		orgID := gid.New(gid.NewTenantID(), 0)
+
+		u, err := c.InitiateWithState(
+			context.Background(),
+			OAuth2State{OrganizationID: orgID.String(), Provider: "TEST"},
+			InitiateOptions{
+				Scopes:               []string{"read:user"},
+				IncludeGrantedScopes: true, // caller requested, but provider does not support
+			},
+			nil,
+		)
+		require.NoError(t, err)
+
+		parsed, err := url.Parse(u)
+		require.NoError(t, err)
+		assert.Equal(t, "consent", parsed.Query().Get("prompt"), "prompt=consent must not be skipped for providers that do not support incremental auth")
+	})
+}
+
+// TestCompleteWithState_ScopeFallback verifies that when the provider's
+// token endpoint returns a successful token response that omits the
+// `scope` field (which RFC 6749 §5.1 allows when the granted scope is
+// identical to the requested scope), CompleteWithState falls back to
+// the RequestedScopes carried in the OAuth2State so the persisted
+// connection still carries the scope set. This is load-bearing for the
+// scope-union logic on subsequent reconnects -- without it we would
+// store empty scope and lose the diff.
+func TestCompleteWithState_ScopeFallback(t *testing.T) {
+	t.Parallel()
+
+	// Fake provider token endpoint: returns a valid token response
+	// with NO `scope` field, matching RFC 6749 §5.1 "identical to
+	// requested" shape.
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, http.MethodPost, r.Method)
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"access_token":"live-token","token_type":"Bearer","expires_in":3600}`))
+	}))
+	defer server.Close()
+
+	c := &OAuth2Connector{
+		ClientID:     "id",
+		ClientSecret: "secret",
+		RedirectURI:  "https://example.com/cb",
+		AuthURL:      "https://provider.example.com/authorize",
+		TokenURL:     server.URL,
+	}
+
+	orgID := gid.New(gid.NewTenantID(), 0)
+	stateData := OAuth2State{
+		OrganizationID:  orgID.String(),
+		Provider:        "TEST",
+		RequestedScopes: []string{"read:user", "write:user"},
+	}
+	stateToken, err := statelesstoken.NewToken(c.ClientSecret, OAuth2TokenType, OAuth2TokenTTL, stateData)
+	require.NoError(t, err)
+
+	// Fabricate a callback request with a code + the signed state.
+	req := httptest.NewRequest(http.MethodGet, "https://example.com/cb?code=the-code&state="+stateToken, nil)
+
+	conn, returnedState, err := c.CompleteWithState(context.Background(), req)
+	require.NoError(t, err)
+	require.NotNil(t, conn)
+	require.NotNil(t, returnedState)
+
+	oauth2Conn, ok := conn.(*OAuth2Connection)
+	require.True(t, ok, "expected *OAuth2Connection, got %T", conn)
+
+	assert.Equal(t, "live-token", oauth2Conn.AccessToken)
+	// The provider omitted scope, so CompleteWithState must fall back
+	// to the RequestedScopes carried in the state token, formatted as
+	// a space-separated RFC 6749 §3.3 scope string (sorted).
+	assert.Equal(t, "read:user write:user", oauth2Conn.Scope)
+	assert.Equal(t, []string{"read:user", "write:user"}, returnedState.RequestedScopes)
 }
--- a/pkg/connector/providers.go
+++ b/pkg/connector/providers.go
@@ -26,10 +26,11 @@ type (
 	// initiate time via InitiateOptions, since the same provider may be used
 	// in multiple contexts requiring different scope sets.
 	providerDefinition struct {
-		AuthURL           string
-		TokenURL          string
-		ExtraAuthParams   map[string]string
-		TokenEndpointAuth string // "post-form" (default), "basic-form", or "basic-json"
+		AuthURL                 string
+		TokenURL                string
+		ExtraAuthParams         map[string]string
+		TokenEndpointAuth       string // "post-form" (default), "basic-form", or "basic-json"
+		SupportsIncrementalAuth bool
 	}
 )

@@ -79,6 +80,7 @@ var (
 				"access_type": "offline",
 				"prompt":      "consent",
 			},
+			SupportsIncrementalAuth: true,
 		},
 		"LINEAR": {
 			AuthURL:  "https://linear.app/oauth/authorize",
@@ -98,5 +100,6 @@ func ApplyProviderDefaults(provider string, redirectURI string, c *OAuth2Connect
 		c.TokenURL = def.TokenURL
 		c.ExtraAuthParams = def.ExtraAuthParams
 		c.TokenEndpointAuth = def.TokenEndpointAuth
+		c.SupportsIncrementalAuth = def.SupportsIncrementalAuth
 	}
 }