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

)
coder · Emyrk · Oct 5, 2023 · Oct 5, 2023 · Oct 5, 2023 · Oct 5, 2023
commit 54509e534deffeaad2342b51e1d1437bd7e6d9d4
diff --git a/cli/server.go b/cli/server.go
@@ -2251,6 +2251,8 @@ func parseExternalAuthProvidersFromEnv(prefix string, environ []string) ([]coder
 			provider.NoRefresh = b
 		case "SCOPES":
 			provider.Scopes = strings.Split(v.Value, " ")
+		case "EXTRA_TOKEN_KEYS":
+			provider.ExtraTokenKeys = strings.Split(v.Value, " ")
 		case "APP_INSTALL_URL":
 			provider.AppInstallURL = v.Value
 		case "APP_INSTALLATIONS_URL":

diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/coderdtest/oidctest/idp.go b/coderd/coderdtest/oidctest/idp.go
@@ -68,6 +68,7 @@ type FakeIDP struct {
 	// "Authorized Redirect URLs". This can be used to emulate that.
 	hookValidRedirectURL func(redirectURL string) error
 	hookUserInfo         func(email string) (jwt.MapClaims, error)
+	hookMutateToken      func(token map[string]interface{})
 	fakeCoderd           func(req *http.Request) (*http.Response, error)
 	hookOnRefresh        func(email string) error
 	// Custom authentication for the client. This is useful if you want
@@ -112,6 +113,14 @@ func WithRefresh(hook func(email string) error) func(*FakeIDP) {
 	}
 }
 
+// WithExtra returns extra fields that be accessed on the returned Oauth Token.
+// These extra fields can override the default fields (id_token, access_token, etc).
+func WithMutateToken(mutateToken func(token map[string]interface{})) func(*FakeIDP) {
+	return func(f *FakeIDP) {
+		f.hookMutateToken = mutateToken
+	}
+}
+
 func WithCustomClientAuth(hook func(t testing.TB, req *http.Request) (url.Values, error)) func(*FakeIDP) {
 	return func(f *FakeIDP) {
 		f.hookAuthenticateClient = hook
@@ -621,6 +630,9 @@ func (f *FakeIDP) httpHandler(t testing.TB) http.Handler {
 			"expires_in":    int64((time.Minute * 5).Seconds()),
 			"id_token":      f.encodeClaims(t, claims),
 		}
+		if f.hookMutateToken != nil {
+			f.hookMutateToken(token)
+		}
 		// Store the claims for the next refresh
 		f.refreshIDTokenClaims.Store(refreshToken, claims)
 

diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go
@@ -4325,6 +4325,7 @@ func (q *FakeQuerier) InsertExternalAuthLink(_ context.Context, arg database.Ins
 		OAuthRefreshToken:      arg.OAuthRefreshToken,
 		OAuthRefreshTokenKeyID: arg.OAuthRefreshTokenKeyID,
 		OAuthExpiry:            arg.OAuthExpiry,
+		OAuthExtra:             arg.OAuthExtra,
 	}
 	q.externalAuthLinks = append(q.externalAuthLinks, gitAuthLink)
 	return gitAuthLink, nil
@@ -5398,6 +5399,7 @@ func (q *FakeQuerier) UpdateExternalAuthLink(_ context.Context, arg database.Upd
 		gitAuthLink.OAuthRefreshToken = arg.OAuthRefreshToken
 		gitAuthLink.OAuthRefreshTokenKeyID = arg.OAuthRefreshTokenKeyID
 		gitAuthLink.OAuthExpiry = arg.OAuthExpiry
+		gitAuthLink.OAuthExtra = arg.OAuthExtra
 		q.externalAuthLinks[index] = gitAuthLink
 
 		return gitAuthLink, nil

diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -514,6 +514,7 @@ func UserLink(t testing.TB, db database.Store, orig database.UserLink) database.
 }
 
 func ExternalAuthLink(t testing.TB, db database.Store, orig database.ExternalAuthLink) database.ExternalAuthLink {
+	msg := takeFirst(&orig.OAuthExtra, &pqtype.NullRawMessage{})
 	link, err := db.InsertExternalAuthLink(genCtx, database.InsertExternalAuthLinkParams{
 		ProviderID:             takeFirst(orig.ProviderID, uuid.New().String()),
 		UserID:                 takeFirst(orig.UserID, uuid.New()),
@@ -524,6 +525,7 @@ func ExternalAuthLink(t testing.TB, db database.Store, orig database.ExternalAut
 		OAuthExpiry:            takeFirst(orig.OAuthExpiry, dbtime.Now().Add(time.Hour*24)),
 		CreatedAt:              takeFirst(orig.CreatedAt, dbtime.Now()),
 		UpdatedAt:              takeFirst(orig.UpdatedAt, dbtime.Now()),
+		OAuthExtra:             *msg,
 	})
 
 	require.NoError(t, err, "insert external auth link")

diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
diff --git a/coderd/database/migrations/000163_external_auth_extra.down.sql b/coderd/database/migrations/000163_external_auth_extra.down.sql
@@ -0,0 +1 @@
+ALTER TABLE external_auth_links DROP COLUMN "oauth_extra";
diff --git a/coderd/database/migrations/000163_external_auth_extra.up.sql b/coderd/database/migrations/000163_external_auth_extra.up.sql
@@ -0,0 +1 @@
+ALTER TABLE external_auth_links ADD COLUMN "oauth_extra" jsonb;
diff --git a/coderd/database/models.go b/coderd/database/models.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/externalauth.sql b/coderd/database/queries/externalauth.sql
@@ -14,7 +14,8 @@ INSERT INTO external_auth_links (
     oauth_access_token_key_id,
     oauth_refresh_token,
     oauth_refresh_token_key_id,
-    oauth_expiry
+    oauth_expiry,
+	oauth_extra
 ) VALUES (
     $1,
     $2,
@@ -24,7 +25,8 @@ INSERT INTO external_auth_links (
     $6,
     $7,
     $8,
-    $9
+    $9,
+	$10
 ) RETURNING *;
 
 -- name: UpdateExternalAuthLink :one
@@ -34,5 +36,6 @@ UPDATE external_auth_links SET
     oauth_access_token_key_id = $5,
     oauth_refresh_token = $6,
     oauth_refresh_token_key_id = $7,
-    oauth_expiry = $8
+    oauth_expiry = $8,
+	oauth_extra = $9
 WHERE provider_id = $1 AND user_id = $2 RETURNING *;
diff --git a/coderd/database/sqlc.yaml b/coderd/database/sqlc.yaml
@@ -53,6 +53,7 @@ overrides:
       oauth_id_token: OAuthIDToken
       oauth_refresh_token: OAuthRefreshToken
       oauth_refresh_token_key_id: OAuthRefreshTokenKeyID
+      oauth_extra: OAuthExtra
       parameter_type_system_hcl: ParameterTypeSystemHCL
       userstatus: UserStatus
       gitsshkey: GitSSHKey

diff --git a/coderd/externalauth.go b/coderd/externalauth.go
@@ -14,6 +14,7 @@ import (
 	"github.com/coder/coder/v2/coderd/httpapi"
 	"github.com/coder/coder/v2/coderd/httpmw"
 	"github.com/coder/coder/v2/codersdk"
+	"github.com/sqlc-dev/pqtype"
 )
 
 // @Summary Get external auth by ID
@@ -132,6 +133,8 @@ func (api *API) postExternalAuthDeviceByID(rw http.ResponseWriter, r *http.Reque
 			OAuthRefreshToken:      token.RefreshToken,
 			OAuthRefreshTokenKeyID: sql.NullString{}, // dbcrypt will set as required
 			OAuthExpiry:            token.Expiry,
+			// No extra data from device auth!
+			OAuthExtra: pqtype.NullRawMessage{},
 		})
 		if err != nil {
 			httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
@@ -150,6 +153,7 @@ func (api *API) postExternalAuthDeviceByID(rw http.ResponseWriter, r *http.Reque
 			OAuthRefreshToken:      token.RefreshToken,
 			OAuthRefreshTokenKeyID: sql.NullString{}, // dbcrypt will update as required
 			OAuthExpiry:            token.Expiry,
+			OAuthExtra:             pqtype.NullRawMessage{},
 		})
 		if err != nil {
 			httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
@@ -201,7 +205,15 @@ func (api *API) externalAuthCallback(externalAuthConfig *externalauth.Config) ht
 			apiKey = httpmw.APIKey(r)
 		)
 
-		_, err := api.Database.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{
+		extra, err := externalAuthConfig.GenerateTokenExtra(state.Token)
+		if err != nil {
+			httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+				Message: "Failed to generate token extra.",
+				Detail:  err.Error(),
+			})
+			return
+		}
+		_, err = api.Database.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{
 			ProviderID: externalAuthConfig.ID,
 			UserID:     apiKey.UserID,
 		})
@@ -224,6 +236,7 @@ func (api *API) externalAuthCallback(externalAuthConfig *externalauth.Config) ht
 				OAuthRefreshToken:      state.Token.RefreshToken,
 				OAuthRefreshTokenKeyID: sql.NullString{}, // dbcrypt will set as required
 				OAuthExpiry:            state.Token.Expiry,
+				OAuthExtra:             extra,
 			})
 			if err != nil {
 				httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
@@ -242,6 +255,7 @@ func (api *API) externalAuthCallback(externalAuthConfig *externalauth.Config) ht
 				OAuthRefreshToken:      state.Token.RefreshToken,
 				OAuthRefreshTokenKeyID: sql.NullString{}, // dbcrypt will update as required
 				OAuthExpiry:            state.Token.Expiry,
+				OAuthExtra:             extra,
 			})
 			if err != nil {
 				httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		ALTER TABLE external_auth_links DROP COLUMN "oauth_extra";
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		ALTER TABLE external_auth_links ADD COLUMN "oauth_extra" jsonb;