chore: prematurely refresh oidc token near expiry during workspace build (#22502)

Closes https://github.com/coder/coder/issues/22429

coderd/provisionerdserver/provisionerdserver.go

@@ -564,7 +564,7 @@ func (s *server) acquireProtoJob(ctx context.Context, job database.ProvisionerJo
 		// The check `s.OIDCConfig != nil` is not as strict, since it can be an interface
 		// pointing to a typed nil.
 		if !reflect.ValueOf(s.OIDCConfig).IsNil() {
-			workspaceOwnerOIDCAccessToken, err = obtainOIDCAccessToken(ctx, s.Database, s.OIDCConfig, owner.ID)
+			workspaceOwnerOIDCAccessToken, err = obtainOIDCAccessToken(ctx, s.Logger, s.Database, s.OIDCConfig, owner.ID)
 			if err != nil {
 				return nil, failJob(fmt.Sprintf("obtain OIDC access token: %s", err))
 			}
@@ -3075,9 +3075,33 @@ func deleteSessionTokenForUserAndWorkspace(ctx context.Context, db database.Stor
 	return nil
 }
 
+func shouldRefreshOIDCToken(link database.UserLink) bool {
+	if link.OAuthRefreshToken == "" {
+		// We cannot refresh even if we wanted to
+		return false
+	}
+
+	if link.OAuthExpiry.IsZero() {
+		// 0 expire means the token never expires, so we shouldn't refresh
+		return false
+	}
+
+	// This handles an edge case where the token is about to expire. A workspace
+	// build takes a non-trivial amount of time. If the token is to expire during the
+	// build, then the build risks failure. To mitigate this, refresh the token
+	// prematurely.
+	//
+	// If an OIDC provider issues short-lived tokens less than our defined period,
+	// the token will always be refreshed on every workspace build.
+	assumeExpiredAt := dbtime.Now().Add(-1 * time.Minute * 10)
+
+	// Return if the token is assumed to be expired.
+	return link.OAuthExpiry.Before(assumeExpiredAt)
+}
+
 // obtainOIDCAccessToken returns a valid OpenID Connect access token
 // for the user if it's able to obtain one, otherwise it returns an empty string.
-func obtainOIDCAccessToken(ctx context.Context, db database.Store, oidcConfig promoauth.OAuth2Config, userID uuid.UUID) (string, error) {
+func obtainOIDCAccessToken(ctx context.Context, logger slog.Logger, db database.Store, oidcConfig promoauth.OAuth2Config, userID uuid.UUID) (string, error) {
 	link, err := db.GetUserLinkByUserIDLoginType(ctx, database.GetUserLinkByUserIDLoginTypeParams{
 		UserID:    userID,
 		LoginType: database.LoginTypeOIDC,
@@ -3089,7 +3113,7 @@ func obtainOIDCAccessToken(ctx context.Context, db database.Store, oidcConfig pr
 		return "", xerrors.Errorf("get owner oidc link: %w", err)
 	}
 
-	if link.OAuthExpiry.Before(dbtime.Now()) && !link.OAuthExpiry.IsZero() && link.OAuthRefreshToken != "" {
+	if shouldRefreshOIDCToken(link) {
 		token, err := oidcConfig.TokenSource(ctx, &oauth2.Token{
 			AccessToken:  link.OAuthAccessToken,
 			RefreshToken: link.OAuthRefreshToken,
@@ -3118,6 +3142,7 @@ func obtainOIDCAccessToken(ctx context.Context, db database.Store, oidcConfig pr
 		if err != nil {
 			return "", xerrors.Errorf("update user link: %w", err)
 		}
+		logger.Info(ctx, "refreshed expired OIDC token for user during workspace build", slog.F("user_id", userID))
 	}
 
 	return link.OAuthAccessToken, nil

coderd/provisionerdserver/provisionerdserver_internal_test.go

@@ -16,13 +16,67 @@ import (
 	"github.com/coder/coder/v2/testutil"
 )
 
+func TestShouldRefreshOIDCToken(t *testing.T) {
+	t.Parallel()
+
+	now := dbtime.Now()
+	testCases := []struct {
+		name string
+		link database.UserLink
+		want bool
+	}{
+		{
+			name: "NoRefreshToken",
+			link: database.UserLink{OAuthExpiry: now.Add(-time.Hour)},
+			want: false,
+		},
+		{
+			name: "ZeroExpiry",
+			link: database.UserLink{OAuthRefreshToken: "refresh"},
+			want: false,
+		},
+		{
+			name: "ExpiredBeyondAssumedWindow",
+			link: database.UserLink{
+				OAuthRefreshToken: "refresh",
+				OAuthExpiry:       now.Add(-20 * time.Minute),
+			},
+			want: true,
+		},
+		{
+			name: "ExpiredWithinAssumedWindow",
+			link: database.UserLink{
+				OAuthRefreshToken: "refresh",
+				OAuthExpiry:       now.Add(-5 * time.Minute),
+			},
+			want: false,
+		},
+		{
+			name: "NotExpired",
+			link: database.UserLink{
+				OAuthRefreshToken: "refresh",
+				OAuthExpiry:       now.Add(time.Hour),
+			},
+			want: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			require.Equal(t, tc.want, shouldRefreshOIDCToken(tc.link))
+		})
+	}
+}
+
 func TestObtainOIDCAccessToken(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
 	t.Run("NoToken", func(t *testing.T) {
 		t.Parallel()
 		db, _ := dbtestutil.NewDB(t)
-		_, err := obtainOIDCAccessToken(ctx, db, nil, uuid.Nil)
+		_, err := obtainOIDCAccessToken(ctx, testutil.Logger(t), db, nil, uuid.Nil)
 		require.NoError(t, err)
 	})
 	t.Run("InvalidConfig", func(t *testing.T) {
@@ -35,7 +89,7 @@ func TestObtainOIDCAccessToken(t *testing.T) {
 			LoginType:   database.LoginTypeOIDC,
 			OAuthExpiry: dbtime.Now().Add(-time.Hour),
 		})
-		_, err := obtainOIDCAccessToken(ctx, db, &oauth2.Config{}, user.ID)
+		_, err := obtainOIDCAccessToken(ctx, testutil.Logger(t), db, &oauth2.Config{}, user.ID)
 		require.NoError(t, err)
 	})
 	t.Run("MissingLink", func(t *testing.T) {
@@ -44,7 +98,7 @@ func TestObtainOIDCAccessToken(t *testing.T) {
 		user := dbgen.User(t, db, database.User{
 			LoginType: database.LoginTypeOIDC,
 		})
-		tok, err := obtainOIDCAccessToken(ctx, db, &oauth2.Config{}, user.ID)
+		tok, err := obtainOIDCAccessToken(ctx, testutil.Logger(t), db, &oauth2.Config{}, user.ID)
 		require.Empty(t, tok)
 		require.NoError(t, err)
 	})
@@ -57,7 +111,7 @@ func TestObtainOIDCAccessToken(t *testing.T) {
 			LoginType:   database.LoginTypeOIDC,
 			OAuthExpiry: dbtime.Now().Add(-time.Hour),
 		})
-		_, err := obtainOIDCAccessToken(ctx, db, &testutil.OAuth2Config{
+		_, err := obtainOIDCAccessToken(ctx, testutil.Logger(t), db, &testutil.OAuth2Config{
 			Token: &oauth2.Token{
 				AccessToken: "token",
 			},