coder · coadler · Dec 3, 2024 · Dec 2, 2024 · Dec 2, 2024 · Dec 2, 2024
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -3330,13 +3330,6 @@ func (q *querier) RegisterWorkspaceProxy(ctx context.Context, arg database.Regis
 	return updateWithReturn(q.log, q.auth, fetch, q.db.RegisterWorkspaceProxy)(ctx, arg)
 }
 
-func (q *querier) RemoveRefreshToken(ctx context.Context, arg database.RemoveRefreshTokenParams) error {
-	fetch := func(ctx context.Context, arg database.RemoveRefreshTokenParams) (database.ExternalAuthLink, error) {
-		return q.db.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{UserID: arg.UserID, ProviderID: arg.ProviderID})
-	}
-	return fetchAndExec(q.log, q.auth, policy.ActionUpdatePersonal, fetch, q.db.RemoveRefreshToken)(ctx, arg)
-}
-
 func (q *querier) RemoveUserFromAllGroups(ctx context.Context, userID uuid.UUID) error {
 	// This is a system function to clear user groups in group sync.
 	if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
@@ -3435,6 +3428,13 @@ func (q *querier) UpdateExternalAuthLink(ctx context.Context, arg database.Updat
 	return fetchAndQuery(q.log, q.auth, policy.ActionUpdatePersonal, fetch, q.db.UpdateExternalAuthLink)(ctx, arg)
 }
 
+func (q *querier) UpdateExternalAuthLinkRefreshToken(ctx context.Context, arg database.UpdateExternalAuthLinkRefreshTokenParams) error {
+	fetch := func(ctx context.Context, arg database.UpdateExternalAuthLinkRefreshTokenParams) (database.ExternalAuthLink, error) {
+		return q.db.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{UserID: arg.UserID, ProviderID: arg.ProviderID})
+	}
+	return fetchAndExec(q.log, q.auth, policy.ActionUpdatePersonal, fetch, q.db.UpdateExternalAuthLinkRefreshToken)(ctx, arg)
+}
+
 func (q *querier) UpdateGitSSHKey(ctx context.Context, arg database.UpdateGitSSHKeyParams) (database.GitSSHKey, error) {
 	fetch := func(ctx context.Context, arg database.UpdateGitSSHKeyParams) (database.GitSSHKey, error) {
 		return q.db.GetGitSSHKey(ctx, arg.UserID)

diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -1282,12 +1282,14 @@ func (s *MethodTestSuite) TestUser() {
 			UserID:     u.ID,
 		}).Asserts(u, policy.ActionUpdatePersonal)
 	}))
-	s.Run("RemoveRefreshToken", s.Subtest(func(db database.Store, check *expects) {
+	s.Run("UpdateExternalAuthLinkRefreshToken", s.Subtest(func(db database.Store, check *expects) {
 		link := dbgen.ExternalAuthLink(s.T(), db, database.ExternalAuthLink{})
-		check.Args(database.RemoveRefreshTokenParams{
-			ProviderID: link.ProviderID,
-			UserID:     link.UserID,
-			UpdatedAt:  link.UpdatedAt,
+		check.Args(database.UpdateExternalAuthLinkRefreshTokenParams{
+			OAuthRefreshToken:      "",
+			OAuthRefreshTokenKeyID: "",
+			ProviderID:             link.ProviderID,
+			UserID:                 link.UserID,
+			UpdatedAt:              link.UpdatedAt,
 		}).Asserts(rbac.ResourceUserObject(link.UserID), policy.ActionUpdatePersonal)
 	}))
 	s.Run("UpdateExternalAuthLink", s.Subtest(func(db database.Store, check *expects) {

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -8556,29 +8556,6 @@ func (q *FakeQuerier) RegisterWorkspaceProxy(_ context.Context, arg database.Reg
 	return database.WorkspaceProxy{}, sql.ErrNoRows
 }
 
-func (q *FakeQuerier) RemoveRefreshToken(_ context.Context, arg database.RemoveRefreshTokenParams) error {
-	if err := validateDatabaseType(arg); err != nil {
-		return err
-	}
-
-	q.mutex.Lock()
-	defer q.mutex.Unlock()
-	for index, gitAuthLink := range q.externalAuthLinks {
-		if gitAuthLink.ProviderID != arg.ProviderID {
-			continue
-		}
-		if gitAuthLink.UserID != arg.UserID {
-			continue
-		}
-		gitAuthLink.UpdatedAt = arg.UpdatedAt
-		gitAuthLink.OAuthRefreshToken = ""
-		q.externalAuthLinks[index] = gitAuthLink
-
-		return nil
-	}
-	return sql.ErrNoRows
-}
-
 func (q *FakeQuerier) RemoveUserFromAllGroups(_ context.Context, userID uuid.UUID) error {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
@@ -8798,6 +8775,29 @@ func (q *FakeQuerier) UpdateExternalAuthLink(_ context.Context, arg database.Upd
 	return database.ExternalAuthLink{}, sql.ErrNoRows
 }
 
+func (q *FakeQuerier) UpdateExternalAuthLinkRefreshToken(_ context.Context, arg database.UpdateExternalAuthLinkRefreshTokenParams) error {
+	if err := validateDatabaseType(arg); err != nil {
+		return err
+	}
+
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+	for index, gitAuthLink := range q.externalAuthLinks {
+		if gitAuthLink.ProviderID != arg.ProviderID {
+			continue
+		}
+		if gitAuthLink.UserID != arg.UserID {
+			continue
+		}
+		gitAuthLink.UpdatedAt = arg.UpdatedAt
+		gitAuthLink.OAuthRefreshToken = arg.OAuthRefreshToken
+		q.externalAuthLinks[index] = gitAuthLink
+
+		return nil
+	}
+	return sql.ErrNoRows
+}
+
 func (q *FakeQuerier) UpdateGitSSHKey(_ context.Context, arg database.UpdateGitSSHKeyParams) (database.GitSSHKey, error) {
 	if err := validateDatabaseType(arg); err != nil {
 		return database.GitSSHKey{}, err

diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.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
@@ -43,13 +43,16 @@ UPDATE external_auth_links SET
 	oauth_extra = $9
 WHERE provider_id = $1 AND user_id = $2 RETURNING *;
 
--- name: RemoveRefreshToken :exec
--- Removing the refresh token disables the refresh behavior for a given
--- auth token. If a refresh token is marked invalid, it is better to remove it
--- then continually attempt to refresh the token.
+-- name: UpdateExternalAuthLinkRefreshToken :exec
 UPDATE
 	external_auth_links
 SET
-	oauth_refresh_token = '',
+	oauth_refresh_token = @oauth_refresh_token,
 	updated_at = @updated_at
-WHERE provider_id = @provider_id AND user_id = @user_id;
+WHERE
+    provider_id = @provider_id
+AND
+    user_id = @user_id
+AND
+    -- Required for sqlc to generate a parameter for the oauth_refresh_token_key_id
+    @oauth_refresh_token_key_id :: text = @oauth_refresh_token_key_id :: text;
diff --git a/coderd/externalauth/externalauth.go b/coderd/externalauth/externalauth.go
@@ -143,10 +143,12 @@ func (c *Config) RefreshToken(ctx context.Context, db database.Store, externalAu
 		// get rid of it. Keeping it around will cause additional refresh
 		// attempts that will fail and cost us api rate limits.
 		if isFailedRefresh(existingToken, err) {
-			dbExecErr := db.RemoveRefreshToken(ctx, database.RemoveRefreshTokenParams{
-				UpdatedAt:  dbtime.Now(),
-				ProviderID: externalAuthLink.ProviderID,
-				UserID:     externalAuthLink.UserID,
+			dbExecErr := db.UpdateExternalAuthLinkRefreshToken(ctx, database.UpdateExternalAuthLinkRefreshTokenParams{
+				OAuthRefreshToken:      "", // It is better to clear the refresh token than to keep retrying.
+				OAuthRefreshTokenKeyID: externalAuthLink.OAuthRefreshTokenKeyID.String,
+				UpdatedAt:              dbtime.Now(),
+				ProviderID:             externalAuthLink.ProviderID,
+				UserID:                 externalAuthLink.UserID,
 			})
 			if dbExecErr != nil {
 				// This error should be rare.

diff --git a/coderd/externalauth/externalauth_test.go b/coderd/externalauth/externalauth_test.go
@@ -190,7 +190,7 @@ func TestRefreshToken(t *testing.T) {
 
 		// Try again with a bad refresh token error
 		// Expect DB call to remove the refresh token
-		mDB.EXPECT().RemoveRefreshToken(gomock.Any(), gomock.Any()).Return(nil).Times(1)
+		mDB.EXPECT().UpdateExternalAuthLinkRefreshToken(gomock.Any(), gomock.Any()).Return(nil).Times(1)
 		refreshErr = &oauth2.RetrieveError{ // github error
 			Response: &http.Response{
 				StatusCode: http.StatusOK,

diff --git a/enterprise/dbcrypt/cipher_internal_test.go b/enterprise/dbcrypt/cipher_internal_test.go
@@ -3,6 +3,8 @@ package dbcrypt
 import (
 	"bytes"
 	"encoding/base64"
+	"os"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -89,3 +91,35 @@ func TestCiphersBackwardCompatibility(t *testing.T) {
 	require.NoError(t, err, "decryption should succeed")
 	require.Equal(t, msg, string(decrypted), "decrypted message should match original message")
 }
+
+// If you're looking here, you're probably in trouble.
+// Here's what you need to do:
+//  1. Get the current CODER_EXTERNAL_TOKEN_ENCRYPTION_KEYS environment variable.
+//  2. Run the following command:
+//     ENCRYPT_ME="<value to encrypt>" CODER_EXTERNAL_TOKEN_ENCRYPTION_KEYS="<secret keys here>" go test -v -count=1 ./enterprise/dbcrypt -test.run='^TestHelpMeEncryptSomeValue$'
+//  3. Copy the value from the test output and do what you need with it.
+func TestHelpMeEncryptSomeValue(t *testing.T) {
+	t.Parallel()
+	t.Skip("this only exists if you need to encrypt a value with dbcrypt, it does not actually test anything")
+
+	valueToEncrypt := os.Getenv("ENCRYPT_ME")
+	t.Logf("valueToEncrypt: %q", valueToEncrypt)
+	keys := os.Getenv("CODER_EXTERNAL_TOKEN_ENCRYPTION_KEYS")
+	require.NotEmpty(t, keys, "Set the CODER_EXTERNAL_TOKEN_ENCRYPTION_KEYS environment variable to use this")
+
+	base64Keys := strings.Split(keys, ",")
+	activeKey := base64Keys[0]
+
+	decodedKey, err := base64.StdEncoding.DecodeString(activeKey)
+	require.NoError(t, err, "the active key should be valid base64")
+
+	cipher, err := cipherAES256(decodedKey)
+	require.NoError(t, err)
+
+	t.Logf("cipher digest: %+v", cipher.HexDigest())
+
+	encryptedEmptyString, err := cipher.Encrypt([]byte(valueToEncrypt))
+	require.NoError(t, err)
+
+	t.Logf("encrypted and base64-encoded: %q", base64.StdEncoding.EncodeToString(encryptedEmptyString))
+}
diff --git a/enterprise/dbcrypt/dbcrypt.go b/enterprise/dbcrypt/dbcrypt.go
@@ -261,6 +261,21 @@ func (db *dbCrypt) UpdateExternalAuthLink(ctx context.Context, params database.U
 	return link, nil
 }
 
+func (db *dbCrypt) UpdateExternalAuthLinkRefreshToken(ctx context.Context, params database.UpdateExternalAuthLinkRefreshTokenParams) error {
+	// We would normally use a sql.NullString here, but sqlc does not want to make
+	// a params struct with a nullable string.
+	var digest sql.NullString
+	if params.OAuthRefreshTokenKeyID != "" {
+		digest.String = params.OAuthRefreshTokenKeyID
+		digest.Valid = true
+	}
+	if err := db.encryptField(&params.OAuthRefreshToken, &digest); err != nil {
+		return err
+	}
+
+	return db.Store.UpdateExternalAuthLinkRefreshToken(ctx, params)
+}
+
 func (db *dbCrypt) GetCryptoKeys(ctx context.Context) ([]database.CryptoKey, error) {
 	keys, err := db.Store.GetCryptoKeys(ctx)
 	if err != nil {