Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion apps/sim/app/api/auth/oauth/utils.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
vi.mock('@/lib/oauth/oauth', () => ({
refreshOAuthToken: vi.fn(),
OAUTH_PROVIDERS: {},
TOKEN_REFRESH_TIMEOUT_MS: 15_000,
}))

const { mockDecryptSecret } = vi.hoisted(() => ({ mockDecryptSecret: vi.fn() }))
Expand Down Expand Up @@ -65,7 +66,9 @@ function mockSelectChain(limitResult: unknown[]) {
* Returns a nested chain: update() -> set() -> where()
*/
function mockUpdateChain() {
const mockWhere = vi.fn().mockResolvedValue({})
/** The rotated write returns the row it matched; an empty result means the chain moved first. */
const mockReturning = vi.fn().mockResolvedValue([{ id: 'account-1' }])
const mockWhere = vi.fn().mockReturnValue({ returning: mockReturning })
const mockSet = vi.fn().mockReturnValue({ where: mockWhere })
mockDb.update.mockReturnValueOnce({ set: mockSet })
return { mockSet, mockWhere }
Expand Down
83 changes: 83 additions & 0 deletions apps/sim/lib/oauth/credential-service.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ vi.mock('@/lib/oauth/microsoft', () => ({
vi.mock('@/lib/oauth/oauth', () => ({
OAUTH_PROVIDERS: {},
refreshOAuthToken: mocks.refreshOAuthToken,
TOKEN_REFRESH_TIMEOUT_MS: 15_000,
}))

vi.mock('@/lib/oauth/quickbooks-client-config', () => ({
Expand Down Expand Up @@ -86,6 +87,7 @@ import {
import { isInstagramProvider, shouldProactivelyRefreshInstagramToken } from '@/lib/oauth/instagram'
import { isMicrosoftProvider } from '@/lib/oauth/microsoft'
import { fanOutSlackTokenChain } from '@/lib/oauth/slack'
import { isTerminalRefreshError, markCredentialDead } from '@/lib/oauth/terminal-errors'
import { GOOGLE_SERVICE_ACCOUNT_PROVIDER_ID } from '@/lib/oauth/types'

const RAW_CREDENTIAL_ID = 'credential-raw-secret-id'
Expand Down Expand Up @@ -243,6 +245,7 @@ describe('resolveCredentialTokenBundle selector privacy', () => {
environment: 'sandbox',
webhookVerifierToken: 'verifier-token',
})
dbChainMockFns.returning.mockResolvedValue([{ id: RAW_ACCOUNT_ID }])
mocks.refreshOAuthToken.mockResolvedValue({
ok: true,
accessToken: 'new-access-token',
Expand Down Expand Up @@ -416,10 +419,13 @@ describe('OAuth access-token refresh headroom', () => {
refreshToken: 'rotated-refresh-token',
expiresIn: 3600,
})
/** The rotated write matches the row unless a test makes the chain move first. */
dbChainMockFns.returning.mockResolvedValue([{ id: RAW_ACCOUNT_ID }])
})

afterEach(() => {
vi.useRealTimers()
vi.mocked(isTerminalRefreshError).mockReturnValue(false)
vi.mocked(isInstagramProvider).mockReturnValue(false)
vi.mocked(shouldProactivelyRefreshInstagramToken).mockReturnValue(false)
vi.mocked(isMicrosoftProvider).mockReturnValue(false)
Expand Down Expand Up @@ -523,6 +529,83 @@ describe('OAuth access-token refresh headroom', () => {
expect(mocks.refreshOAuthToken).not.toHaveBeenCalled()
})

it('sizes the lease and the follower wait past the provider timeout for every provider', async () => {
queueCredentialAccount(createOAuthAccount())
await resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test')
expect(mocks.withLeaderLock).toHaveBeenCalledWith(
expect.objectContaining({ ttlSec: 30, maxWaitMs: 30_000 })
)
})

it('rotates the chain only from the refresh token the refresh started from', async () => {
queueCredentialAccount(createOAuthAccount())
await expect(
resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test')
).resolves.toEqual({ accessToken: 'refreshed-access-token' })
expect(dbChainMockFns.set).toHaveBeenCalledWith(
expect.objectContaining({ refreshToken: 'rotated-refresh-token' })
)
const guard = JSON.stringify(dbChainMockFns.where.mock.calls.at(-1))
expect(guard).toContain('account.refreshToken')
expect(guard).toContain('original-refresh-token')
})

it('returns no token when the rotation write finds the account gone', async () => {
queueCredentialAccount(createOAuthAccount())
dbChainMockFns.returning.mockResolvedValueOnce([])
queueTableRows(account, [])
await expect(
resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test')
).resolves.toBeNull()
expect(mocks.logger.warn).toHaveBeenCalledWith(
'Rotation write found no account; the credential is gone',
expect.anything()
)
})

it('does not flag a credential dead when a terminal failure follows a newer rotation', async () => {
queueCredentialAccount(createOAuthAccount())
vi.mocked(isTerminalRefreshError).mockReturnValue(true)
mocks.refreshOAuthToken.mockResolvedValue({ ok: false, errorCode: 'invalid_grant' })
queueTableRows(account, [
{
...createOAuthAccount(3_600_000),
accessToken: 'winner-token',
refreshToken: 'winner-refresh-token',
},
])
await expect(
resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test')
).resolves.toEqual({ accessToken: 'winner-token' })
expect(markCredentialDead).not.toHaveBeenCalled()
})

it('flags a credential dead on a terminal failure when its chain did not move', async () => {
queueCredentialAccount(createOAuthAccount())
vi.mocked(isTerminalRefreshError).mockReturnValue(true)
mocks.refreshOAuthToken.mockResolvedValue({ ok: false, errorCode: 'invalid_grant' })
queueTableRows(account, [createOAuthAccount()])
await expect(
resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test')
).resolves.toBeNull()
expect(markCredentialDead).toHaveBeenCalledWith(expect.any(String), 'invalid_grant')
})

it('uses the stored chain when the rotation write loses to a newer one', async () => {
queueCredentialAccount(createOAuthAccount())
/** Another writer rotated first: no row still holds the token this refresh started from. */
dbChainMockFns.returning.mockResolvedValueOnce([])
queueTableRows(account, [{ ...createOAuthAccount(3_600_000), accessToken: 'winner-token' }])
await expect(
resolveCredentialTokenBundle(RAW_CREDENTIAL_ID, RAW_USER_ID, 'test')
).resolves.toEqual({ accessToken: 'winner-token' })
expect(mocks.refreshOAuthToken).toHaveBeenCalledTimes(1)
expect(mocks.logger.warn).toHaveBeenCalledWith(
'Rotation write lost to a newer chain; using the stored token',
expect.anything()
)
})

it.each([
{ remainingMs: 6_000, refresh: true },
{ remainingMs: 300_000, refresh: true },
Expand Down
110 changes: 86 additions & 24 deletions apps/sim/lib/oauth/credential-service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ import {
isMicrosoftProvider,
PROACTIVE_REFRESH_THRESHOLD_DAYS,
} from '@/lib/oauth/microsoft'
import { refreshOAuthToken } from '@/lib/oauth/oauth'
import { refreshOAuthToken, TOKEN_REFRESH_TIMEOUT_MS } from '@/lib/oauth/oauth'
import { decryptQuickBooksOAuthClientConfig } from '@/lib/oauth/quickbooks-client-config'
import { getOAuthRefreshCoordinationIdentity } from '@/lib/oauth/refresh-coordination'
import {
Expand Down Expand Up @@ -872,18 +872,50 @@ function isOAuthAccessTokenExpiring(
}

/**
* Slack lock budgets sized past `TOKEN_REFRESH_TIMEOUT_MS` (15s) in
* lib/oauth/oauth.ts: installation-keyed locks make every sibling row's request
* a follower of one refresh, so the TTL covers the provider call plus generous
* headroom for the surrounding DB reads and the fan-out write, and followers
* poll for the lock's full lifetime so a slow-but-successful refresh is still
* observed rather than reported as a failure. These budgets are latency knobs,
* not correctness guarantees — chain integrity under lock expiry or unlocked
* concurrent writers is enforced by the version-guarded fan-out
* (`ifChainUnchangedSince` in lib/oauth/slack.ts).
* Lock budgets sized past the provider call: the lease covers
* {@link TOKEN_REFRESH_TIMEOUT_MS} plus headroom for the account read before it and
* the rotated write after it, so a leader still talking to a slow provider keeps its
* lease instead of letting a second leader start a competing rotation; and followers
* poll for the lease's full lifetime, so a slow-but-successful refresh is observed
* rather than reported as a failure. Both are latency knobs, not correctness
* guarantees: a lease is only ever a lease, and chain integrity under lock expiry or
* an unlocked writer is enforced at the write, which rotates a chain only from the
* refresh token it started from (`ifChainUnchangedSince` for a Slack installation).
*/
const SLACK_LOCK_TTL_SEC = 30
const SLACK_FOLLOWER_MAX_WAIT_MS = SLACK_LOCK_TTL_SEC * 1000
const REFRESH_LOCK_HEADROOM_MS = 15_000
const REFRESH_LOCK_TTL_SEC = Math.ceil((TOKEN_REFRESH_TIMEOUT_MS + REFRESH_LOCK_HEADROOM_MS) / 1000)
const REFRESH_FOLLOWER_MAX_WAIT_MS = REFRESH_LOCK_TTL_SEC * 1000

interface StoredChain {
accessToken: string | null
accessTokenExpiresAt: Date | null
refreshToken: string | null
}

/** The chain an account row holds now, or nothing when the account is gone. */
async function readStoredChain(accountId: string): Promise<StoredChain | undefined> {
const [stored] = await db
.select({
accessToken: account.accessToken,
accessTokenExpiresAt: account.accessTokenExpiresAt,
refreshToken: account.refreshToken,
})
.from(account)
.where(eq(account.id, accountId))
.limit(1)
return stored
}

/**
* The stored access token when it can still serve a request, as a follower would take it: a
* chain another writer just rotated carries one, and a token that has already expired is no
* answer at all.
*/
function usableStoredToken(stored: StoredChain, providerId: string): string | null {
return stored.accessToken && !isOAuthAccessTokenExpiring(stored.accessTokenExpiresAt, providerId)
? stored.accessToken
: null
}

async function performCoalescedRefresh({
accountId,
Expand Down Expand Up @@ -926,11 +958,8 @@ async function performCoalescedRefresh({
const refreshPromise = coalesceLocally(lockKey, () =>
withLeaderLock<string>({
key: lockKey,
// Installation-keyed Slack locks gather followers from every sibling row,
// so their wait and the lock TTL must outlast the 15s provider timeout —
// the 3s/10s defaults would fail followers early and let a second leader
// start a concurrent rotation mid-refresh.
...(slackTeamId ? { maxWaitMs: SLACK_FOLLOWER_MAX_WAIT_MS, ttlSec: SLACK_LOCK_TTL_SEC } : {}),
ttlSec: REFRESH_LOCK_TTL_SEC,
maxWaitMs: REFRESH_FOLLOWER_MAX_WAIT_MS,
onLeader: async () => {
try {
let refreshTokenToUse = refreshToken
Expand Down Expand Up @@ -981,18 +1010,25 @@ async function performCoalescedRefresh({
message: result.message,
})
if (result.errorCode && isTerminalRefreshError(result.errorCode)) {
// A refresh that lost a race with a concurrent connect fails with
// a revoked/rotated-out token even though the installation just
// got a live chain — dead-flagging then would take down a healthy
// credential for an hour.
// A refresh that lost a race with a concurrent connect or a newer
// rotation fails with a revoked/rotated-out token even though the
// account just got a live chain — dead-flagging then would take
// down a healthy credential for an hour.
if (
slackChainVersion &&
(await hasSlackChainMoved(slackTeamId!, slackChainVersion))
) {
logger.info('Skipping dead flag: Slack chain moved during refresh', logContext)
} else {
await markCredentialDead(scopeKey, result.errorCode)
return null
}
if (!slackTeamId) {
const stored = await readStoredChain(accountId)
if (stored && stored.refreshToken !== refreshToken) {
logger.info('Skipping dead flag: chain moved during refresh', logContext)
return usableStoredToken(stored, providerId)
}
}
await markCredentialDead(scopeKey, result.errorCode)
}
return null
}
Expand Down Expand Up @@ -1027,7 +1063,33 @@ async function performCoalescedRefresh({
)
}

await db.update(account).set(updateData).where(eq(account.id, accountId))
/**
* The chain is rotated only from the refresh token this refresh started from.
* A lease is not mutual exclusion: it can expire under a slow provider or a
* paused process while the leader is still running, and an unconditional
* write would then let this refresh overwrite a newer rotation with a chain
* the provider has already retired, which the next refresh pays for as
* `invalid_grant` and, under reuse detection, as a revoked grant. When no row
* matches, another writer rotated first: its chain is the live one, so this
* caller uses what is stored and never retries the provider.
*/
const rotated = await db
.update(account)
.set(updateData)
.where(and(eq(account.id, accountId), eq(account.refreshToken, refreshToken)))
Comment thread
waleedlatif1 marked this conversation as resolved.
.returning({ id: account.id })
if (rotated.length === 0) {
const stored = await readStoredChain(accountId)
if (!stored) {
logger.warn('Rotation write found no account; the credential is gone', logContext)
return null
}
logger.warn(
'Rotation write lost to a newer chain; using the stored token',
logContext
)
return usableStoredToken(stored, providerId)
}
}

logger.info('Successfully refreshed access token', logContext)
Expand Down
6 changes: 5 additions & 1 deletion apps/sim/lib/oauth/github-installation-credential.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,11 @@ vi.mock('@/lib/oauth/github-installation', () => ({
parseGitHubInstallationBinding: mocks.parseBinding,
resolveGitHubInstallationAccessToken: mocks.resolveToken,
}))
vi.mock('@/lib/oauth/oauth', () => ({ OAUTH_PROVIDERS: {}, refreshOAuthToken: vi.fn() }))
vi.mock('@/lib/oauth/oauth', () => ({
OAUTH_PROVIDERS: {},
refreshOAuthToken: vi.fn(),
TOKEN_REFRESH_TIMEOUT_MS: 15_000,
}))

import { resolveServiceAccountToken } from '@/lib/oauth/credential-service'
import { GITHUB_INSTALLATION_PROVIDER_ID } from '@/lib/oauth/github-installation-types'
Expand Down
2 changes: 1 addition & 1 deletion apps/sim/lib/oauth/oauth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2250,7 +2250,7 @@ function safeOAuthErrorCode(value: unknown, secrets: string[]): string | undefin
* Without this bound a hung endpoint would wedge every joiner on that key until
* the undici socket defaults (~5 min) gave up.
*/
const TOKEN_REFRESH_TIMEOUT_MS = 15_000
export const TOKEN_REFRESH_TIMEOUT_MS = 15_000

function parseOAuthResponse(responseText: string): unknown {
try {
Expand Down
Loading