mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-05 14:52:09 +03:00
fix(oauth): per-connection mutex for rotating refresh tokens (#1885)
Integrated into release/v3.7.9
This commit is contained in:
@@ -13,6 +13,11 @@ const CACHE_SECRET = "omniroute-token-cache";
|
||||
// Key: "provider:sha256(refreshToken)" → Value: Promise<result>
|
||||
const refreshPromiseCache = new Map();
|
||||
|
||||
// Per-connection mutex: prevents parallel OAuth refresh for rotating tokens.
|
||||
// Key: connectionId → Value: { promise, waiters }
|
||||
// Primary dedup when credentials.connectionId is present; refreshPromiseCache is fallback.
|
||||
const connectionRefreshMutex = new Map();
|
||||
|
||||
type RefreshLogger = {
|
||||
info?: (tag: string, message: string, data?: Record<string, unknown>) => void;
|
||||
warn?: (tag: string, message: string, data?: Record<string, unknown>) => void;
|
||||
@@ -415,8 +420,8 @@ export async function refreshQwenToken(refreshToken, log, proxyConfig: unknown =
|
||||
/**
|
||||
* Specialized refresh for Codex (OpenAI) OAuth tokens.
|
||||
* OpenAI uses rotating (one-time-use) refresh tokens.
|
||||
* Returns { error: 'refresh_token_reused' } when the token has already been consumed,
|
||||
* so callers can stop retrying and request re-authentication.
|
||||
* Returns { error: 'unrecoverable_refresh_error', code } when the token has already been
|
||||
* consumed or is invalid, so callers can stop retrying and request re-authentication.
|
||||
*/
|
||||
export async function refreshCodexToken(refreshToken, log, proxyConfig: unknown = null) {
|
||||
try {
|
||||
@@ -831,9 +836,18 @@ export function isUnrecoverableRefreshError(result) {
|
||||
|
||||
/**
|
||||
* Get access token for a specific provider (with deduplication).
|
||||
* If a refresh is already in-flight for the same provider+token,
|
||||
* subsequent calls share the existing promise instead of making
|
||||
* parallel OAuth requests.
|
||||
*
|
||||
* Deduplication strategy (two layers):
|
||||
* 1. Per-connection mutex (primary): if credentials.connectionId is present, all concurrent
|
||||
* callers for that connection share one in-flight promise regardless of which token they
|
||||
* loaded. This prevents refresh_token_reused errors with rotating (one-time-use) tokens,
|
||||
* e.g. Codex/OpenAI, where callers that loaded credentials at different times may hold
|
||||
* different token strings but refer to the same connection.
|
||||
* 2. Token-hash fallback: if no connectionId, dedup by provider+sha256(refreshToken) as before.
|
||||
*
|
||||
* Additionally, when connectionId is present, the stale-token check reads the DB to detect
|
||||
* whether another process already refreshed the token. If the DB token is still valid it is
|
||||
* returned immediately without a new upstream call.
|
||||
*/
|
||||
export async function getAccessToken(provider, credentials, log, proxyConfig: unknown = null) {
|
||||
if (!credentials || !credentials.refreshToken || typeof credentials.refreshToken !== "string") {
|
||||
@@ -841,6 +855,52 @@ export async function getAccessToken(provider, credentials, log, proxyConfig: un
|
||||
return null;
|
||||
}
|
||||
|
||||
const connectionId = credentials.connectionId;
|
||||
|
||||
// ── Layer 1: per-connection mutex ──────────────────────────────────────────
|
||||
if (connectionId && typeof connectionId === "string") {
|
||||
const existing = connectionRefreshMutex.get(connectionId);
|
||||
if (existing) {
|
||||
existing.waiters++;
|
||||
log?.info?.("TOKEN_REFRESH", "Concurrent refresh detected — sharing in-flight refresh", {
|
||||
provider,
|
||||
connectionId,
|
||||
waiters: existing.waiters,
|
||||
});
|
||||
return existing.promise;
|
||||
}
|
||||
|
||||
const entry = { promise: null, waiters: 0 };
|
||||
entry.promise = _getAccessTokenWithStalenessCheck(provider, credentials, log, proxyConfig).finally(() => {
|
||||
connectionRefreshMutex.delete(connectionId);
|
||||
});
|
||||
connectionRefreshMutex.set(connectionId, entry);
|
||||
return entry.promise;
|
||||
}
|
||||
|
||||
// ── Layer 2: token-hash fallback (no connectionId) ─────────────────────────
|
||||
const cacheKey = getRefreshCacheKey(provider, credentials.refreshToken);
|
||||
|
||||
if (refreshPromiseCache.has(cacheKey)) {
|
||||
log?.info?.("TOKEN_REFRESH", `Reusing in-flight refresh for ${provider}`);
|
||||
return refreshPromiseCache.get(cacheKey);
|
||||
}
|
||||
|
||||
const refreshPromise = _getAccessTokenInternal(provider, credentials, log, proxyConfig).finally(
|
||||
() => {
|
||||
refreshPromiseCache.delete(cacheKey);
|
||||
}
|
||||
);
|
||||
|
||||
refreshPromiseCache.set(cacheKey, refreshPromise);
|
||||
return refreshPromise;
|
||||
}
|
||||
|
||||
/**
|
||||
* Internal helper: performs the DB staleness check then calls the actual refresh.
|
||||
* Only called from the per-connection mutex path (Layer 1 above).
|
||||
*/
|
||||
async function _getAccessTokenWithStalenessCheck(provider, credentials, log, proxyConfig) {
|
||||
// RACE CONDITION PREVENTION:
|
||||
// If the credentials object in memory is stale (e.g. it waited in a semaphore while another
|
||||
// request refreshed the token), using its OLD refreshToken will cause the provider (e.g. OpenAI)
|
||||
@@ -886,23 +946,7 @@ export async function getAccessToken(provider, credentials, log, proxyConfig: un
|
||||
}
|
||||
}
|
||||
|
||||
const cacheKey = getRefreshCacheKey(provider, credentials.refreshToken);
|
||||
|
||||
// If a refresh is already in-flight, reuse it
|
||||
if (refreshPromiseCache.has(cacheKey)) {
|
||||
log?.info?.("TOKEN_REFRESH", `Reusing in-flight refresh for ${provider}`);
|
||||
return refreshPromiseCache.get(cacheKey);
|
||||
}
|
||||
|
||||
// Start a new refresh and cache the promise
|
||||
const refreshPromise = _getAccessTokenInternal(provider, credentials, log, proxyConfig).finally(
|
||||
() => {
|
||||
refreshPromiseCache.delete(cacheKey);
|
||||
}
|
||||
);
|
||||
|
||||
refreshPromiseCache.set(cacheKey, refreshPromise);
|
||||
return refreshPromise;
|
||||
return _getAccessTokenInternal(provider, credentials, log, proxyConfig);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -1036,6 +1080,18 @@ export function isProviderBlocked(provider: string): boolean {
|
||||
return false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get active per-connection mutex entries (for diagnostics/metrics).
|
||||
* Returns a snapshot of connections that have an in-flight refresh and their waiter count.
|
||||
*/
|
||||
export function getConnectionRefreshMutexStatus(): Record<string, { waiters: number }> {
|
||||
const result: Record<string, { waiters: number }> = {};
|
||||
for (const [connectionId, entry] of connectionRefreshMutex.entries()) {
|
||||
result[connectionId] = { waiters: entry.waiters };
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get circuit breaker status for all providers (for diagnostics).
|
||||
*/
|
||||
|
||||
@@ -17,6 +17,32 @@ import {
|
||||
getAllAccessTokens as _getAllAccessTokens,
|
||||
} from "@omniroute/open-sse/services/tokenRefresh.ts";
|
||||
|
||||
// Per-connection mutex: prevents concurrent OAuth refresh for rotating tokens.
|
||||
// Key = connectionId, Value = { promise: in-flight refresh, waiters: count of callers sharing it }
|
||||
const connectionRefreshMutex = new Map<string, { promise: Promise<any>; waiters: number }>();
|
||||
|
||||
export async function withConnectionRefreshMutex<T>(
|
||||
connectionId: string,
|
||||
fn: () => Promise<T>
|
||||
): Promise<T> {
|
||||
const existing = connectionRefreshMutex.get(connectionId);
|
||||
if (existing) {
|
||||
existing.waiters++;
|
||||
log.info("TOKEN_REFRESH", "Concurrent refresh detected — sharing in-flight refresh", {
|
||||
connectionId,
|
||||
waiters: existing.waiters,
|
||||
});
|
||||
return existing.promise as Promise<T>;
|
||||
}
|
||||
|
||||
const entry: { promise: Promise<T>; waiters: number } = { promise: null as any, waiters: 0 };
|
||||
entry.promise = fn().finally(() => {
|
||||
connectionRefreshMutex.delete(connectionId);
|
||||
});
|
||||
connectionRefreshMutex.set(connectionId, entry);
|
||||
return entry.promise;
|
||||
}
|
||||
|
||||
export const TOKEN_EXPIRY_BUFFER_MS = BUFFER_MS;
|
||||
|
||||
export const refreshAccessToken = async (
|
||||
@@ -146,7 +172,12 @@ export async function checkAndRefreshToken(provider: string, credentials: any) {
|
||||
expiresIn: Math.round((expiresAt - now) / 1000),
|
||||
});
|
||||
|
||||
const newCredentials = await getAccessToken(provider, updatedCredentials);
|
||||
const connectionId: string | undefined = updatedCredentials.connectionId;
|
||||
const newCredentials = connectionId
|
||||
? await withConnectionRefreshMutex(connectionId, () =>
|
||||
getAccessToken(provider, updatedCredentials)
|
||||
)
|
||||
: await getAccessToken(provider, updatedCredentials);
|
||||
if (newCredentials && newCredentials.accessToken) {
|
||||
await updateProviderCredentials(updatedCredentials.connectionId, newCredentials);
|
||||
|
||||
|
||||
@@ -24,6 +24,7 @@ const {
|
||||
getAllAccessTokens,
|
||||
isProviderBlocked,
|
||||
getCircuitBreakerStatus,
|
||||
getConnectionRefreshMutexStatus,
|
||||
refreshWithRetry,
|
||||
} = tokenRefresh;
|
||||
|
||||
@@ -907,3 +908,200 @@ test("isProviderBlocked clears expired circuit-breaker entries once cooldown pas
|
||||
assert.equal(getCircuitBreakerStatus()[provider], undefined);
|
||||
});
|
||||
});
|
||||
|
||||
// ─── Per-connection mutex tests ────────────────────────────────────────────────
|
||||
|
||||
test("getAccessToken per-connection mutex: 5 concurrent callers fire exactly one upstream call", async () => {
|
||||
const log = createLog();
|
||||
let upstreamCallCount = 0;
|
||||
|
||||
await withPatchedProperties(
|
||||
PROVIDERS,
|
||||
{ "custom-oauth-conn-mutex": { tokenUrl: "https://auth.example.com/token" } },
|
||||
async () => {
|
||||
await withMockedFetch(
|
||||
async () => {
|
||||
upstreamCallCount++;
|
||||
// Simulate 50ms upstream latency so all 5 callers are concurrent
|
||||
await new Promise((r) => setTimeout(r, 50));
|
||||
return jsonResponse({
|
||||
access_token: "new-access-token",
|
||||
refresh_token: "new-refresh-token",
|
||||
expires_in: 3600,
|
||||
});
|
||||
},
|
||||
async () => {
|
||||
const credentials = {
|
||||
connectionId: "conn-abc-123",
|
||||
refreshToken: "old-refresh-token",
|
||||
};
|
||||
|
||||
const results = await Promise.all([
|
||||
getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log),
|
||||
getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log),
|
||||
getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log),
|
||||
getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log),
|
||||
getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log),
|
||||
]);
|
||||
|
||||
assert.equal(upstreamCallCount, 1, "upstream called exactly once");
|
||||
for (const result of results) {
|
||||
assert.equal(result?.accessToken, "new-access-token", "all callers got same token");
|
||||
assert.equal(result?.refreshToken, "new-refresh-token");
|
||||
}
|
||||
// All results are the same object reference (shared promise)
|
||||
assert.strictEqual(results[0], results[1]);
|
||||
assert.strictEqual(results[1], results[4]);
|
||||
}
|
||||
);
|
||||
}
|
||||
);
|
||||
});
|
||||
|
||||
test("getAccessToken per-connection mutex: logs concurrent refresh with waiter count", async () => {
|
||||
const log = createLog();
|
||||
|
||||
await withPatchedProperties(
|
||||
PROVIDERS,
|
||||
{ "custom-oauth-conn-mutex": { tokenUrl: "https://auth.example.com/token" } },
|
||||
async () => {
|
||||
await withMockedFetch(
|
||||
async () => {
|
||||
await new Promise((r) => setTimeout(r, 20));
|
||||
return jsonResponse({ access_token: "tok", refresh_token: "rtok", expires_in: 600 });
|
||||
},
|
||||
async () => {
|
||||
const credentials = { connectionId: "conn-log-test", refreshToken: "rt" };
|
||||
await Promise.all([
|
||||
getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log),
|
||||
getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log),
|
||||
getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log),
|
||||
]);
|
||||
|
||||
const concurrentLogs = log.entries.filter(
|
||||
(e) =>
|
||||
e.level === "info" &&
|
||||
e.message === "Concurrent refresh detected — sharing in-flight refresh"
|
||||
);
|
||||
assert.ok(concurrentLogs.length >= 1, "logged at least one concurrent refresh event");
|
||||
assert.ok(
|
||||
concurrentLogs.some((e) => e.meta?.connectionId === "conn-log-test"),
|
||||
"log includes connectionId"
|
||||
);
|
||||
assert.ok(
|
||||
concurrentLogs.some((e) => typeof e.meta?.waiters === "number" && e.meta.waiters >= 1),
|
||||
"log includes waiter count"
|
||||
);
|
||||
}
|
||||
);
|
||||
}
|
||||
);
|
||||
});
|
||||
|
||||
test("getAccessToken per-connection mutex: failed refresh propagates null to all waiters (idempotent error)", async () => {
|
||||
const log = createLog();
|
||||
|
||||
await withPatchedProperties(
|
||||
PROVIDERS,
|
||||
{ "custom-oauth-conn-mutex": { tokenUrl: "https://auth.example.com/token" } },
|
||||
async () => {
|
||||
await withMockedFetch(
|
||||
async () => {
|
||||
await new Promise((r) => setTimeout(r, 20));
|
||||
// 400 response causes refreshAccessToken to return null
|
||||
return new Response("bad_request", { status: 400 });
|
||||
},
|
||||
async () => {
|
||||
const credentials = { connectionId: "conn-fail-test", refreshToken: "expired-rt" };
|
||||
const results = await Promise.all([
|
||||
getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log),
|
||||
getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log),
|
||||
getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log),
|
||||
]);
|
||||
|
||||
for (const result of results) {
|
||||
assert.equal(result, null, "failed refresh returns null to all waiters");
|
||||
}
|
||||
// Mutex cleaned up after failure
|
||||
assert.equal(
|
||||
getConnectionRefreshMutexStatus()["conn-fail-test"],
|
||||
undefined,
|
||||
"mutex entry removed after failure"
|
||||
);
|
||||
}
|
||||
);
|
||||
}
|
||||
);
|
||||
});
|
||||
|
||||
test("getAccessToken per-connection mutex: different connections run independently", async () => {
|
||||
const log = createLog();
|
||||
let upstreamCallCount = 0;
|
||||
|
||||
await withPatchedProperties(
|
||||
PROVIDERS,
|
||||
{ "custom-oauth-conn-mutex": { tokenUrl: "https://auth.example.com/token" } },
|
||||
async () => {
|
||||
await withMockedFetch(
|
||||
async () => {
|
||||
upstreamCallCount++;
|
||||
await new Promise((r) => setTimeout(r, 20));
|
||||
return jsonResponse({
|
||||
access_token: `access-${upstreamCallCount}`,
|
||||
refresh_token: `refresh-${upstreamCallCount}`,
|
||||
expires_in: 600,
|
||||
});
|
||||
},
|
||||
async () => {
|
||||
const [groupA, groupB] = await Promise.all([
|
||||
Promise.all([
|
||||
getAccessToken("custom-oauth-conn-mutex", { connectionId: "conn-A", refreshToken: "rt-a" }, log),
|
||||
getAccessToken("custom-oauth-conn-mutex", { connectionId: "conn-A", refreshToken: "rt-a" }, log),
|
||||
]),
|
||||
Promise.all([
|
||||
getAccessToken("custom-oauth-conn-mutex", { connectionId: "conn-B", refreshToken: "rt-b" }, log),
|
||||
getAccessToken("custom-oauth-conn-mutex", { connectionId: "conn-B", refreshToken: "rt-b" }, log),
|
||||
]),
|
||||
]);
|
||||
|
||||
assert.equal(upstreamCallCount, 2, "one upstream call per distinct connection");
|
||||
assert.strictEqual(groupA[0], groupA[1], "conn-A callers share same result");
|
||||
assert.strictEqual(groupB[0], groupB[1], "conn-B callers share same result");
|
||||
assert.notStrictEqual(groupA[0], groupB[0], "conn-A and conn-B got different results");
|
||||
}
|
||||
);
|
||||
}
|
||||
);
|
||||
});
|
||||
|
||||
test("getAccessToken per-connection mutex: mutex cleared after success, next call re-fires upstream", async () => {
|
||||
const log = createLog();
|
||||
let upstreamCallCount = 0;
|
||||
|
||||
await withPatchedProperties(
|
||||
PROVIDERS,
|
||||
{ "custom-oauth-conn-mutex": { tokenUrl: "https://auth.example.com/token" } },
|
||||
async () => {
|
||||
await withMockedFetch(
|
||||
async () => {
|
||||
upstreamCallCount++;
|
||||
return jsonResponse({
|
||||
access_token: `access-${upstreamCallCount}`,
|
||||
refresh_token: `refresh-${upstreamCallCount}`,
|
||||
expires_in: 600,
|
||||
});
|
||||
},
|
||||
async () => {
|
||||
const credentials = { connectionId: "conn-refire", refreshToken: "rt" };
|
||||
|
||||
const first = await getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log);
|
||||
const second = await getAccessToken("custom-oauth-conn-mutex", { ...credentials }, log);
|
||||
|
||||
assert.equal(upstreamCallCount, 2, "each sequential call fires upstream once");
|
||||
assert.equal(first?.accessToken, "access-1");
|
||||
assert.equal(second?.accessToken, "access-2");
|
||||
}
|
||||
);
|
||||
}
|
||||
);
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user