feat: allow storing extra oauth token properties in the database (#10152)

2026-09-24 15:04:27 +08:00 · 2023-10-09 18:49:30 -05:00
parent 35538e1051
commit 863c2e7b64
25 changed files with 223 additions and 60 deletions
@@ -15,6 +15,7 @@ import (
 	"golang.org/x/xerrors"

 	"github.com/google/go-github/v43/github"
+	"github.com/sqlc-dev/pqtype"
 	xgithub "golang.org/x/oauth2/github"

 	"github.com/coder/coder/v2/coderd/database"
@@ -44,6 +45,14 @@ type Config struct {
 	// DisplayIcon is the path to an image that will be displayed to the user.
 	DisplayIcon string

+	// ExtraTokenKeys is a list of extra properties to
+	// store in the database returned from the token endpoint.
+	//
+	// e.g. Slack returns `authed_user` in the token which is
+	// a payload that contains information about the authenticated
+	// user.
+	ExtraTokenKeys []string
+
 	// NoRefresh stops Coder from using the refresh token
 	// to renew the access token.
 	//
@@ -69,6 +78,25 @@ type Config struct {
 	AppInstallationsURL string
 }

+// GenerateTokenExtra generates the extra token data to store in the database.
+func (c *Config) GenerateTokenExtra(token *oauth2.Token) (pqtype.NullRawMessage, error) {
+	if len(c.ExtraTokenKeys) == 0 {
+		return pqtype.NullRawMessage{}, nil
+	}
+	extraMap := map[string]interface{}{}
+	for _, key := range c.ExtraTokenKeys {
+		extraMap[key] = token.Extra(key)
+	}
+	data, err := json.Marshal(extraMap)
+	if err != nil {
+		return pqtype.NullRawMessage{}, err
+	}
+	return pqtype.NullRawMessage{
+		RawMessage: data,
+		Valid:      true,
+	}, nil
+}
+
 // RefreshToken automatically refreshes the token if expired and permitted.
 // It returns the token and a bool indicating if the token is valid.
 func (c *Config) RefreshToken(ctx context.Context, db database.Store, externalAuthLink database.ExternalAuthLink) (database.ExternalAuthLink, bool, error) {
@@ -101,6 +129,12 @@ func (c *Config) RefreshToken(ctx context.Context, db database.Store, externalAu
 		// we aren't trying to surface an error, we're just trying to obtain a valid token.
 		return externalAuthLink, false, nil
 	}
+
+	extra, err := c.GenerateTokenExtra(token)
+	if err != nil {
+		return externalAuthLink, false, xerrors.Errorf("generate token extra: %w", err)
+	}
+
 	r := retry.New(50*time.Millisecond, 200*time.Millisecond)
 	// See the comment below why the retry and cancel is required.
 	retryCtx, retryCtxCancel := context.WithTimeout(ctx, time.Second)
@@ -135,6 +169,7 @@ validate:
 			OAuthRefreshToken:      token.RefreshToken,
 			OAuthRefreshTokenKeyID: sql.NullString{}, // dbcrypt will update as required
 			OAuthExpiry:            token.Expiry,
+			OAuthExtra:             extra,
 		})
 		if err != nil {
 			return updatedAuthLink, false, xerrors.Errorf("update external auth link: %w", err)
@@ -539,6 +574,14 @@ var defaults = map[codersdk.EnhancedExternalAuthProvider]codersdk.ExternalAuthCo
 		DeviceCodeURL:       "https://github.com/login/device/code",
 		AppInstallationsURL: "https://api.github.com/user/installations",
 	},
+	codersdk.EnhancedExternalAuthProviderSlack: {
+		AuthURL:     "https://slack.com/oauth/v2/authorize",
+		TokenURL:    "https://slack.com/api/oauth.v2.access",
+		DisplayName: "Slack",
+		DisplayIcon: "/icon/slack.svg",
+		// See: https://api.slack.com/authentication/oauth-v2#exchanging
+		ExtraTokenKeys: []string{"authed_user"},
+	},
 }

 // jwtConfig is a new OAuth2 config that uses a custom
@@ -2,6 +2,7 @@ package externalauth_test

 import (
 	"context"
+	"encoding/json"
 	"net/http"
 	"net/url"
 	"testing"
@@ -43,7 +44,7 @@ func TestRefreshToken(t *testing.T) {
 					return nil, xerrors.New("should not be called")
 				}),
 			},
-			GitConfigOpt: func(cfg *externalauth.Config) {
+			ExternalAuthOpt: func(cfg *externalauth.Config) {
 				cfg.NoRefresh = true
 			},
 		})
@@ -74,7 +75,7 @@ func TestRefreshToken(t *testing.T) {
 					return jwt.MapClaims{}, nil
 				}),
 			},
-			GitConfigOpt: func(cfg *externalauth.Config) {
+			ExternalAuthOpt: func(cfg *externalauth.Config) {
 				cfg.NoRefresh = true
 			},
 		})
@@ -117,7 +118,7 @@ func TestRefreshToken(t *testing.T) {
 					return jwt.MapClaims{}, xerrors.New(staticError)
 				}),
 			},
-			GitConfigOpt: func(cfg *externalauth.Config) {
+			ExternalAuthOpt: func(cfg *externalauth.Config) {
 			},
 		})

@@ -142,7 +143,7 @@ func TestRefreshToken(t *testing.T) {
 					return jwt.MapClaims{}, oidctest.StatusError(http.StatusUnauthorized, xerrors.New(staticError))
 				}),
 			},
-			GitConfigOpt: func(cfg *externalauth.Config) {
+			ExternalAuthOpt: func(cfg *externalauth.Config) {
 			},
 		})

@@ -175,7 +176,7 @@ func TestRefreshToken(t *testing.T) {
 					return jwt.MapClaims{}, oidctest.StatusError(http.StatusUnauthorized, xerrors.New(staticError))
 				}),
 			},
-			GitConfigOpt: func(cfg *externalauth.Config) {
+			ExternalAuthOpt: func(cfg *externalauth.Config) {
 				cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String()
 			},
 		})
@@ -205,7 +206,7 @@ func TestRefreshToken(t *testing.T) {
 					return jwt.MapClaims{}, nil
 				}),
 			},
-			GitConfigOpt: func(cfg *externalauth.Config) {
+			ExternalAuthOpt: func(cfg *externalauth.Config) {
 				cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String()
 			},
 		})
@@ -236,7 +237,7 @@ func TestRefreshToken(t *testing.T) {
 					return jwt.MapClaims{}, nil
 				}),
 			},
-			GitConfigOpt: func(cfg *externalauth.Config) {
+			ExternalAuthOpt: func(cfg *externalauth.Config) {
 				cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String()
 			},
 			DB: db,
@@ -260,6 +261,41 @@ func TestRefreshToken(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, updated.OAuthAccessToken, dbLink.OAuthAccessToken, "token is updated in the DB")
 	})
+
+	t.Run("WithExtra", func(t *testing.T) {
+		t.Parallel()
+
+		db := dbfake.New()
+		fake, config, link := setupOauth2Test(t, testConfig{
+			FakeIDPOpts: []oidctest.FakeIDPOpt{
+				oidctest.WithMutateToken(func(token map[string]interface{}) {
+					token["authed_user"] = map[string]interface{}{
+						"access_token": token["access_token"],
+					}
+				}),
+			},
+			ExternalAuthOpt: func(cfg *externalauth.Config) {
+				cfg.Type = codersdk.EnhancedExternalAuthProviderSlack.String()
+				cfg.ExtraTokenKeys = []string{"authed_user"}
+				cfg.ValidateURL = ""
+			},
+			DB: db,
+		})
+
+		ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil))
+		// Force a refresh
+		link.OAuthExpiry = expired
+
+		updated, ok, err := config.RefreshToken(ctx, db, link)
+		require.NoError(t, err)
+		require.True(t, ok)
+		require.True(t, updated.OAuthExtra.Valid)
+		extra := map[string]interface{}{}
+		require.NoError(t, json.Unmarshal(updated.OAuthExtra.RawMessage, &extra))
+		mapping, ok := extra["authed_user"].(map[string]interface{})
+		require.True(t, ok)
+		require.Equal(t, updated.OAuthAccessToken, mapping["access_token"])
+	})
 }

 func TestConvertYAML(t *testing.T) {
@@ -344,7 +380,7 @@ func TestConvertYAML(t *testing.T) {
 type testConfig struct {
 	FakeIDPOpts         []oidctest.FakeIDPOpt
 	CoderOIDCConfigOpts []func(cfg *coderd.OIDCConfig)
-	GitConfigOpt        func(cfg *externalauth.Config)
+	ExternalAuthOpt     func(cfg *externalauth.Config)
 	// If DB is passed in, the link will be inserted into the DB.
 	DB database.Store
 }
@@ -367,7 +403,7 @@ func setupOauth2Test(t *testing.T, settings testConfig) (*oidctest.FakeIDP, *ext
 		ID:           providerID,
 		ValidateURL:  fake.WellknownConfig().UserInfoURL,
 	}
-	settings.GitConfigOpt(config)
+	settings.ExternalAuthOpt(config)

 	oauthToken, err := fake.GenerateAuthenticatedToken(jwt.MapClaims{
 		"email": "test@coder.com",