support multiple token refresh

This commit is contained in:
msukkari 2025-11-03 19:39:12 -08:00
parent e6498531aa
commit 4f3369b62f
5 changed files with 238 additions and 152 deletions

View file

@ -18,7 +18,7 @@ import { hasEntitlement } from '@sourcebot/shared';
import { onCreateUser } from '@/lib/authUtils';
import { getAuditService } from '@/ee/features/audit/factory';
import { SINGLE_TENANT_ORG_ID } from './lib/constants';
import { refreshOAuthToken } from '@/ee/features/permissionSyncing/actions';
import { refreshIntegrationTokens } from '@/ee/features/permissionSyncing/tokenRefresh';
const auditService = getAuditService();
const eeIdentityProviders = hasEntitlement("sso") ? await getEEIdentityProviders() : [];
@ -36,16 +36,19 @@ declare module 'next-auth' {
user: {
id: string;
} & DefaultSession['user'];
integrationProviderErrors?: Record<string, string>;
}
}
declare module 'next-auth/jwt' {
interface JWT {
userId: string;
accessToken?: string;
refreshToken?: string;
expiresAt?: number;
provider?: string;
integrationTokens?: Record<string, {
accessToken: string;
refreshToken: string;
expiresAt: number;
error?: string;
}>;
error?: string;
}
}
@ -193,41 +196,24 @@ export const { handlers, signIn, signOut, auth } = NextAuth({
token.userId = user.id;
}
if (account) {
token.accessToken = account.access_token;
token.refreshToken = account.refresh_token;
token.expiresAt = account.expires_at;
token.provider = account.provider;
// When a user links a new account, store the tokens if it's an integration provider
if (account && hasEntitlement('permission-syncing')) {
if (account.access_token && account.refresh_token && account.expires_at) {
token.integrationTokens = token.integrationTokens || {};
token.integrationTokens[account.provider] = {
accessToken: account.access_token,
refreshToken: account.refresh_token,
expiresAt: account.expires_at,
};
}
}
if (hasEntitlement('permission-syncing') &&
token.provider &&
['github', 'gitlab'].includes(token.provider) &&
token.expiresAt &&
token.refreshToken) {
const now = Math.floor(Date.now() / 1000);
const bufferTimeS = 5 * 60;
if (now >= (token.expiresAt - bufferTimeS)) {
try {
const refreshedTokens = await refreshOAuthToken(
token.provider,
token.refreshToken,
// Refresh all integration provider tokens that are about to expire
if (hasEntitlement('permission-syncing') && token.integrationTokens) {
token.integrationTokens = await refreshIntegrationTokens(
token.integrationTokens,
token.userId
);
if (refreshedTokens) {
token.accessToken = refreshedTokens.accessToken;
token.refreshToken = refreshedTokens.refreshToken ?? token.refreshToken;
token.expiresAt = refreshedTokens.expiresAt;
} else {
token.error = 'RefreshTokenError';
}
} catch (error) {
console.error('Error refreshing token:', error);
token.error = 'RefreshTokenError';
}
}
}
return token;
@ -240,6 +226,18 @@ export const { handlers, signIn, signOut, auth } = NextAuth({
// Propagate the userId to the session.
id: token.userId,
}
// Pass only integration provider errors to the session (not sensitive tokens)
if (token.integrationTokens) {
const errors: Record<string, string> = {};
for (const [provider, tokenData] of Object.entries(token.integrationTokens)) {
if (tokenData.error) {
errors[provider] = tokenData.error;
}
}
if (Object.keys(errors).length > 0) {
session.integrationProviderErrors = errors;
}
}
return session;
},
},

View file

@ -8,9 +8,8 @@ import { env } from "@/env.mjs";
import { OrgRole } from "@sourcebot/db";
import { cookies } from "next/headers";
import { OPTIONAL_PROVIDERS_LINK_SKIPPED_COOKIE_NAME } from "@/lib/constants";
import { getTokenFromConfig } from '@sourcebot/crypto';
import { IntegrationIdentityProviderState } from "@/ee/features/permissionSyncing/types";
import { GitHubIdentityProviderConfig, GitLabIdentityProviderConfig } from "@sourcebot/schemas/v3/index.type";
import { auth } from "@/auth";
const logger = createLogger('web-ee-permission-syncing-actions');
@ -32,6 +31,10 @@ export const getIntegrationProviderStates = async () => sew(() =>
}
});
// Fetch the session to get token errors
const session = await auth();
const providerErrors = session?.integrationProviderErrors;
const integrationProviderState: IntegrationIdentityProviderState[] = [];
for (const integrationProviderConfig of integrationProviderConfigs) {
if (integrationProviderConfig.purpose === "integration") {
@ -41,11 +44,14 @@ export const getIntegrationProviderStates = async () => sew(() =>
const isLinked = !!linkedAccount;
const isRequired = integrationProviderConfig.required ?? true;
const providerError = providerErrors?.[integrationProviderConfig.provider];
integrationProviderState.push({
id: integrationProviderConfig.provider,
required: isRequired,
isLinked,
linkedAccountId: linkedAccount?.providerAccountId
linkedAccountId: linkedAccount?.providerAccountId,
error: providerError
} as IntegrationIdentityProviderState);
}
}
@ -99,90 +105,3 @@ export const skipOptionalProvidersLink = async () => sew(async () => {
return true;
});
export const refreshOAuthToken = async (
provider: string,
refreshToken: string,
userId: string
): Promise<{ accessToken: string; refreshToken: string | null; expiresAt: number } | null> => {
try {
// Load config and find the provider configuration
const config = await loadConfig(env.CONFIG_PATH);
const identityProviders = config?.identityProviders ?? [];
const providerConfig = identityProviders.find(
idp => idp.provider === provider
) as GitHubIdentityProviderConfig | GitLabIdentityProviderConfig;
if (!providerConfig || !('clientId' in providerConfig) || !('clientSecret' in providerConfig)) {
logger.error(`Provider config not found or invalid for: ${provider}`);
return null;
}
// Get client credentials from config
const clientId = await getTokenFromConfig(providerConfig.clientId);
const clientSecret = await getTokenFromConfig(providerConfig.clientSecret);
const baseUrl = 'baseUrl' in providerConfig && providerConfig.baseUrl
? await getTokenFromConfig(providerConfig.baseUrl)
: undefined;
let url: string;
if (baseUrl) {
url = provider === 'github'
? `${baseUrl}/login/oauth/access_token`
: `${baseUrl}/oauth/token`;
} else if (provider === 'github') {
url = 'https://github.com/login/oauth/access_token';
} else if (provider === 'gitlab') {
url = 'https://gitlab.com/oauth/token';
} else {
logger.error(`Unsupported provider for token refresh: ${provider}`);
return null;
}
const response = await fetch(url, {
method: 'POST',
headers: {
'Content-Type': 'application/x-www-form-urlencoded',
'Accept': 'application/json',
},
body: new URLSearchParams({
client_id: clientId,
client_secret: clientSecret,
grant_type: 'refresh_token',
refresh_token: refreshToken,
}),
});
if (!response.ok) {
const errorText = await response.text();
logger.error(`Failed to refresh ${provider} token: ${response.status} ${errorText}`);
return null;
}
const data = await response.json();
const result = {
accessToken: data.access_token,
refreshToken: data.refresh_token ?? null,
expiresAt: data.expires_in ? Math.floor(Date.now() / 1000) + data.expires_in : 0,
};
const { prisma } = await import('@/prisma');
await prisma.account.updateMany({
where: {
userId: userId,
provider: provider,
},
data: {
access_token: result.accessToken,
refresh_token: result.refreshToken,
expires_at: result.expiresAt,
},
});
return result;
} catch (error) {
logger.error(`Error refreshing ${provider} token:`, error);
return null;
}
};

View file

@ -1,5 +1,5 @@
import { getAuthProviderInfo } from "@/lib/utils";
import { Check, X } from "lucide-react";
import { Check, X, AlertCircle } from "lucide-react";
import { Card, CardDescription, CardHeader, CardTitle } from "@/components/ui/card";
import { ProviderIcon } from "./providerIcon";
import { ProviderInfo } from "./providerInfo";
@ -41,6 +41,7 @@ export function IntegrationProviderCard({
/>
</CardTitle>
<CardDescription className="text-xs">
<div className="flex flex-col gap-1.5">
{integrationProviderState.isLinked? (
<div className="flex items-center gap-2">
<div className="flex items-center gap-1.5">
@ -66,6 +67,15 @@ export function IntegrationProviderCard({
</span>
</div>
)}
{integrationProviderState.error && (
<div className="flex items-center gap-1.5">
<AlertCircle className="h-3.5 w-3.5 text-destructive" />
<span className="text-destructive font-medium">
Token refresh failed - please reconnect
</span>
</div>
)}
</div>
</CardDescription>
</div>
</div>

View file

@ -0,0 +1,158 @@
import { loadConfig } from "@sourcebot/shared";
import { env } from "@/env.mjs";
import { createLogger } from "@sourcebot/logger";
import { getTokenFromConfig } from '@sourcebot/crypto';
import { GitHubIdentityProviderConfig, GitLabIdentityProviderConfig } from "@sourcebot/schemas/v3/index.type";
const logger = createLogger('web-ee-token-refresh');
export type IntegrationToken = {
accessToken: string;
refreshToken: string;
expiresAt: number;
error?: string;
};
export type IntegrationTokensMap = Record<string, IntegrationToken>;
export async function refreshIntegrationTokens(
currentTokens: IntegrationTokensMap | undefined,
userId: string
): Promise<IntegrationTokensMap> {
if (!currentTokens) {
return {};
}
const now = Math.floor(Date.now() / 1000);
const bufferTimeS = 5 * 60; // 5 minutes
const updatedTokens: IntegrationTokensMap = { ...currentTokens };
// Refresh tokens for each integration provider
await Promise.all(
Object.entries(currentTokens).map(async ([provider, tokenData]) => {
if (provider !== 'github' && provider !== 'gitlab') {
return;
}
if (tokenData.expiresAt && now >= (tokenData.expiresAt - bufferTimeS)) {
try {
logger.info(`Refreshing token for provider: ${provider}`);
const refreshedTokens = await refreshOAuthToken(
provider,
tokenData.refreshToken,
userId
);
if (refreshedTokens) {
updatedTokens[provider] = {
accessToken: refreshedTokens.accessToken,
refreshToken: refreshedTokens.refreshToken ?? tokenData.refreshToken,
expiresAt: refreshedTokens.expiresAt,
};
logger.info(`Successfully refreshed token for provider: ${provider}`);
} else {
logger.error(`Failed to refresh token for provider: ${provider}`);
updatedTokens[provider] = {
...tokenData,
error: 'RefreshTokenError',
};
}
} catch (error) {
logger.error(`Error refreshing token for provider ${provider}:`, error);
updatedTokens[provider] = {
...tokenData,
error: 'RefreshTokenError',
};
}
}
})
);
return updatedTokens;
}
export async function refreshOAuthToken(
provider: string,
refreshToken: string,
userId: string
): Promise<{ accessToken: string; refreshToken: string | null; expiresAt: number } | null> {
try {
const config = await loadConfig(env.CONFIG_PATH);
const identityProviders = config?.identityProviders ?? [];
const providerConfig = identityProviders.find(idp => idp.provider === provider);
if (!providerConfig) {
logger.error(`Provider config not found or invalid for: ${provider}`);
return null;
}
// Get client credentials from config
const integrationProviderConfig = providerConfig as GitHubIdentityProviderConfig | GitLabIdentityProviderConfig
const clientId = await getTokenFromConfig(integrationProviderConfig.clientId);
const clientSecret = await getTokenFromConfig(integrationProviderConfig.clientSecret);
const baseUrl = 'baseUrl' in integrationProviderConfig && integrationProviderConfig.baseUrl
? await getTokenFromConfig(integrationProviderConfig.baseUrl)
: undefined;
let url: string;
if (baseUrl) {
url = provider === 'github'
? `${baseUrl}/login/oauth/access_token`
: `${baseUrl}/oauth/token`;
} else if (provider === 'github') {
url = 'https://github.com/login/oauth/access_token';
} else if (provider === 'gitlab') {
url = 'https://gitlab.com/oauth/token';
} else {
logger.error(`Unsupported provider for token refresh: ${provider}`);
return null;
}
const response = await fetch(url, {
method: 'POST',
headers: {
'Content-Type': 'application/x-www-form-urlencoded',
'Accept': 'application/json',
},
body: new URLSearchParams({
client_id: clientId,
client_secret: clientSecret,
grant_type: 'refresh_token',
refresh_token: refreshToken,
}),
});
if (!response.ok) {
const errorText = await response.text();
logger.error(`Failed to refresh ${provider} token: ${response.status} ${errorText}`);
return null;
}
const data = await response.json();
const result = {
accessToken: data.access_token,
refreshToken: data.refresh_token ?? null,
expiresAt: data.expires_in ? Math.floor(Date.now() / 1000) + data.expires_in : 0,
};
const { prisma } = await import('@/prisma');
await prisma.account.updateMany({
where: {
userId: userId,
provider: provider,
},
data: {
access_token: result.accessToken,
refresh_token: result.refreshToken,
expires_at: result.expiresAt,
},
});
return result;
} catch (error) {
logger.error(`Error refreshing ${provider} token:`, error);
return null;
}
}

View file

@ -3,4 +3,5 @@ export type IntegrationIdentityProviderState = {
required: boolean;
isLinked: boolean;
linkedAccountId?: string;
error?: string;
};