Skip to content

Commit ee3748e

Browse files
author
kjgbot
committed
Harden integration auth recovery
1 parent e26f86b commit ee3748e

16 files changed

Lines changed: 655 additions & 71 deletions

‎src/main/auth.test.ts‎

Lines changed: 69 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,17 @@ function tokenHash(accessToken: string): string {
9191
return createHash('sha256').update(accessToken).digest('hex')
9292
}
9393

94+
function accountKeyForToken(apiUrl: string, accessToken: string, subject: string): string {
95+
return createHash('sha256')
96+
.update(`${apiUrl}\0token-subject:${subject}`)
97+
.digest('hex')
98+
}
99+
100+
function jwtWithSubject(subject: string, nonce: string): string {
101+
const encode = (value: unknown): string => Buffer.from(JSON.stringify(value), 'utf8').toString('base64url')
102+
return `${encode({ alg: 'none' })}.${encode({ sub: subject, nonce })}.sig`
103+
}
104+
94105
describe('getAccountWorkspaceId', () => {
95106
let userDataDir: string
96107

@@ -138,7 +149,8 @@ describe('getAccountWorkspaceId', () => {
138149
expect(String(calledUrl)).toBe('https://cloud.example/api/v1/auth/whoami')
139150

140151
const meta = readMeta(userDataDir)
141-
expect(meta?.accountWorkspace).toEqual({
152+
expect(meta?.accountWorkspace).toMatchObject({
153+
accountKey: expect.any(String),
142154
tokenHash: tokenHash('cld_at_abc'),
143155
workspaceId: 'ws-from-current'
144156
})
@@ -203,6 +215,40 @@ describe('getAccountWorkspaceId', () => {
203215
const { getAccountWorkspaceId } = await import('./auth')
204216
const id = await getAccountWorkspaceId()
205217

218+
expect(id).toBe('ws-cached')
219+
expect(mock.fetchMock).not.toHaveBeenCalled()
220+
expect(readMeta(userDataDir)?.accountWorkspace).toMatchObject({
221+
accountKey: expect.any(String),
222+
tokenHash: tokenHash('cld_at_cached'),
223+
workspaceId: 'ws-cached'
224+
})
225+
})
226+
227+
it('uses the cached workspace id across token rotation for the same account subject', async () => {
228+
const oldToken = jwtWithSubject('user-1', 'old')
229+
const newToken = jwtWithSubject('user-1', 'new')
230+
writeAuthJson(userDataDir, {
231+
accessToken: newToken,
232+
refreshToken: 'cld_rt_rotated',
233+
apiUrl: 'https://cloud.example'
234+
})
235+
const configDir = join(userDataDir, 'config')
236+
mkdirSync(configDir, { recursive: true })
237+
writeFileSync(
238+
join(configDir, 'auth-meta.json'),
239+
JSON.stringify({
240+
apiUrl: 'https://cloud.example',
241+
accountWorkspace: {
242+
accountKey: accountKeyForToken('https://cloud.example', oldToken, 'user-1'),
243+
tokenHash: tokenHash(oldToken),
244+
workspaceId: 'ws-cached'
245+
}
246+
})
247+
)
248+
249+
const { getAccountWorkspaceId } = await import('./auth')
250+
const id = await getAccountWorkspaceId()
251+
206252
expect(id).toBe('ws-cached')
207253
expect(mock.fetchMock).not.toHaveBeenCalled()
208254
})
@@ -239,7 +285,8 @@ describe('getAccountWorkspaceId', () => {
239285
expect(mock.fetchMock).toHaveBeenCalledTimes(1)
240286

241287
const meta = readMeta(userDataDir)
242-
expect(meta?.accountWorkspace).toEqual({
288+
expect(meta?.accountWorkspace).toMatchObject({
289+
accountKey: expect.any(String),
243290
tokenHash: tokenHash('cld_at_new'),
244291
workspaceId: 'ws-fresh'
245292
})
@@ -288,7 +335,7 @@ describe('getAccountWorkspaceId', () => {
288335
expect(mock.fetchMock).toHaveBeenCalledTimes(2)
289336
})
290337

291-
it('throws account-workspace-required when whoami responds with a non-OK status', async () => {
338+
it('throws cloud-auth-required when whoami rejects the access token', async () => {
292339
writeAuthJson(userDataDir, {
293340
accessToken: 'cld_at_401',
294341
refreshToken: 'cld_rt_401',
@@ -302,7 +349,25 @@ describe('getAccountWorkspaceId', () => {
302349
})
303350

304351
const { getAccountWorkspaceId } = await import('./auth')
305-
await expect(getAccountWorkspaceId()).rejects.toThrowError('account-workspace-required')
352+
await expect(getAccountWorkspaceId()).rejects.toThrowError('cloud-auth-required:whoami-http-401')
353+
})
354+
355+
it('throws account-workspace-required with a failure class when whoami returns a server error', async () => {
356+
writeAuthJson(userDataDir, {
357+
accessToken: 'cld_at_500',
358+
refreshToken: 'cld_rt_500',
359+
apiUrl: 'https://cloud.example'
360+
})
361+
mock.fetchMock.mockResolvedValue({
362+
ok: false,
363+
status: 500,
364+
statusText: 'Internal Server Error',
365+
json: async () => ({ error: 'No active workspace' })
366+
})
367+
368+
const { getAccountWorkspaceId } = await import('./auth')
369+
await expect(getAccountWorkspaceId({ retryAttempts: 1, retryDelayMs: 0 }))
370+
.rejects.toThrowError('account-workspace-required:whoami-http-500')
306371
})
307372
})
308373

‎src/main/auth.ts‎

Lines changed: 75 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ const CLOUD_API_URL = process.env.RELAY_CLOUD_URL || 'https://agentrelay.dev/clo
1818
const TOKEN_EXPIRY_BUFFER_MS = 60_000
1919
const ACCOUNT_WORKSPACE_RETRY_ATTEMPTS = 8
2020
const ACCOUNT_WORKSPACE_RETRY_DELAY_MS = 500
21+
const warnedWhoamiWorkspaceFailures = new Set<string>()
2122

2223
interface AuthStatus {
2324
loggedIn: boolean
@@ -26,7 +27,8 @@ interface AuthStatus {
2627
}
2728

2829
type AccountWorkspaceCache = {
29-
tokenHash: string
30+
accountKey?: string
31+
tokenHash?: string
3032
workspaceId: string
3133
}
3234

@@ -128,16 +130,35 @@ function accountWorkspaceTokenHash(accessToken: string): string {
128130
return createHash('sha256').update(accessToken).digest('hex')
129131
}
130132

133+
function accountWorkspaceCacheMatches(
134+
cached: AccountWorkspaceCache | undefined,
135+
auth: Pick<CloudAuth, 'accountKey' | 'accessToken'>
136+
): boolean {
137+
if (!cached?.workspaceId.trim()) return false
138+
return cached.accountKey === auth.accountKey ||
139+
cached.tokenHash === accountWorkspaceTokenHash(auth.accessToken)
140+
}
141+
131142
function delay(ms: number): Promise<void> {
132143
return new Promise((resolve) => setTimeout(resolve, ms))
133144
}
134145

135146
function saveAuthMeta(tokens: Pick<StoredTokens, 'apiUrl' | 'user'> & Partial<Pick<StoredTokens, 'accessToken'>>): void {
136147
const previous = loadAuthMeta()
148+
const accountKey = tokens.accessToken
149+
? deriveCloudAuthAccountKey(tokens.apiUrl, tokens.accessToken, tokens.user)
150+
: undefined
137151
const tokenHash = tokens.accessToken ? accountWorkspaceTokenHash(tokens.accessToken) : undefined
138152
const accountWorkspace =
139-
tokenHash && previous.accountWorkspace?.tokenHash === tokenHash
140-
? previous.accountWorkspace
153+
accountKey && accountWorkspaceCacheMatches(previous.accountWorkspace, {
154+
accountKey,
155+
accessToken: tokens.accessToken || ''
156+
})
157+
? {
158+
...previous.accountWorkspace,
159+
accountKey,
160+
...(tokenHash ? { tokenHash } : {})
161+
}
141162
: undefined
142163
const meta = {
143164
apiUrl: tokens.apiUrl,
@@ -240,7 +261,21 @@ function clearTokens(): void {
240261
}
241262
}
242263

243-
async function fetchWhoamiPayload(apiUrl: string, accessToken: string): Promise<unknown | undefined> {
264+
type WhoamiPayloadResult =
265+
| { ok: true; data: unknown }
266+
| { ok: false; failureClass: string; status?: number }
267+
268+
function whoamiFailureClassForStatus(status: number): string {
269+
return `whoami-http-${status}`
270+
}
271+
272+
function warnWhoamiWorkspaceFailure(failureClass: string): void {
273+
if (warnedWhoamiWorkspaceFailures.has(failureClass)) return
274+
warnedWhoamiWorkspaceFailures.add(failureClass)
275+
console.warn('[auth] Account workspace whoami lookup failed:', { failureClass })
276+
}
277+
278+
async function fetchWhoamiPayload(apiUrl: string, accessToken: string): Promise<WhoamiPayloadResult> {
244279
const controller = new AbortController()
245280
const timeout = setTimeout(() => controller.abort(), 2500)
246281

@@ -249,10 +284,21 @@ async function fetchWhoamiPayload(apiUrl: string, accessToken: string): Promise<
249284
headers: { Authorization: `Bearer ${accessToken}` },
250285
signal: controller.signal
251286
})
252-
if (!res.ok) return undefined
253-
return await res.json() as unknown
254-
} catch {
255-
return undefined
287+
if (!res.ok) {
288+
return {
289+
ok: false,
290+
status: res.status,
291+
failureClass: whoamiFailureClassForStatus(res.status)
292+
}
293+
}
294+
return { ok: true, data: await res.json() as unknown }
295+
} catch (error) {
296+
return {
297+
ok: false,
298+
failureClass: error instanceof Error && error.name === 'AbortError'
299+
? 'whoami-timeout'
300+
: 'whoami-network'
301+
}
256302
} finally {
257303
clearTimeout(timeout)
258304
}
@@ -277,6 +323,7 @@ function saveAccountWorkspaceCache(auth: CloudAuth, workspaceId: string): void {
277323
apiUrl: auth.apiUrl || previous.apiUrl?.trim() || CLOUD_API_URL,
278324
user: previous.user,
279325
accountWorkspace: {
326+
accountKey: auth.accountKey,
280327
tokenHash: accountWorkspaceTokenHash(auth.accessToken),
281328
workspaceId
282329
}
@@ -286,7 +333,9 @@ function saveAccountWorkspaceCache(auth: CloudAuth, workspaceId: string): void {
286333

287334
async function fetchWhoami(apiUrl: string, accessToken: string): Promise<UserInfo | undefined> {
288335
try {
289-
const data = await fetchWhoamiPayload(apiUrl, accessToken)
336+
const payload = await fetchWhoamiPayload(apiUrl, accessToken)
337+
if (!payload.ok) return undefined
338+
const data = payload.data
290339
const record = isRecord(data) ? data : {}
291340
const userRecord = firstObject(record, ['user']) || record
292341
const organizationRecord = firstObject(record, ['organization', 'org'])
@@ -646,24 +695,36 @@ export async function getAccountWorkspaceId(options: AccountWorkspaceIdOptions =
646695
const auth = await resolveCloudAuth()
647696
if (!auth) throw new Error('cloud-auth-required')
648697

649-
const tokenHash = accountWorkspaceTokenHash(auth.accessToken)
650698
const cached = loadAuthMeta().accountWorkspace
651-
if (cached?.tokenHash === tokenHash && cached.workspaceId.trim()) {
699+
if (accountWorkspaceCacheMatches(cached, auth)) {
652700
return cached.workspaceId.trim()
653701
}
654702

655703
const retryAttempts = Math.max(1, Math.floor(options.retryAttempts ?? 1))
656704
const retryDelayMs = Math.max(0, Math.floor(options.retryDelayMs ?? ACCOUNT_WORKSPACE_RETRY_DELAY_MS))
657705
let workspaceId: string | undefined
706+
let failureClass = 'whoami-no-workspace-in-payload'
658707

659708
for (let attempt = 1; attempt <= retryAttempts; attempt += 1) {
660-
const data = await fetchWhoamiPayload(auth.apiUrl, auth.accessToken)
661-
workspaceId = accountWorkspaceIdFromWhoami(data)
709+
const payload = await fetchWhoamiPayload(auth.apiUrl, auth.accessToken)
710+
if (!payload.ok) {
711+
failureClass = payload.failureClass
712+
if (payload.status === 401 || payload.status === 403) {
713+
warnWhoamiWorkspaceFailure(failureClass)
714+
throw new Error(`cloud-auth-required:${failureClass}`)
715+
}
716+
} else {
717+
workspaceId = accountWorkspaceIdFromWhoami(payload.data)
718+
failureClass = workspaceId ? '' : 'whoami-no-workspace-in-payload'
719+
}
662720
if (workspaceId) break
663721
if (attempt < retryAttempts) await delay(retryDelayMs)
664722
}
665723

666-
if (!workspaceId) throw new Error('account-workspace-required')
724+
if (!workspaceId) {
725+
warnWhoamiWorkspaceFailure(failureClass)
726+
throw new Error(`account-workspace-required:${failureClass}`)
727+
}
667728

668729
saveAccountWorkspaceCache(auth, workspaceId)
669730
return workspaceId

‎src/main/integration-mounts.test.ts‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -449,6 +449,7 @@ describe('IntegrationMountManager', () => {
449449
'utf8'
450450
)
451451
expect(healthObserver).toHaveBeenCalledWith({
452+
type: 'auth-stall',
452453
remotePath: '/slack/channels/C123/messages',
453454
status: 'writeback-pending',
454455
pendingWriteback: 1,
@@ -492,4 +493,68 @@ describe('IntegrationMountManager', () => {
492493
expect(healthObserver).toHaveBeenCalledTimes(1)
493494
expect(mock.mountInputs.filter((input) => input.remotePath === '/slack/channels/C123/messages')).toHaveLength(2)
494495
})
496+
497+
it('rejects and reports cloud auth recovery when startup has no usable tokens', async () => {
498+
mock.currentAuth = null
499+
const manager = new IntegrationMountManager()
500+
const healthObserver = vi.fn()
501+
manager.setHealthObserver(healthObserver)
502+
503+
await expect(manager.ensureMounted([
504+
{
505+
provider: 'slack',
506+
mountPaths: ['/slack/channels/C123/messages']
507+
}
508+
])).rejects.toThrow('cloud-auth-required')
509+
510+
expect(mock.mountInputs).toHaveLength(0)
511+
expect(healthObserver).toHaveBeenCalledWith({
512+
type: 'auth-required',
513+
reason: 'cloud-auth-required',
514+
message: 'cloud-auth-required'
515+
})
516+
})
517+
518+
it('rejects and reports workspace recovery when whoami lacks an account workspace', async () => {
519+
mock.getAccountWorkspaceId.mockRejectedValueOnce(new Error('account-workspace-required'))
520+
const manager = new IntegrationMountManager()
521+
const healthObserver = vi.fn()
522+
manager.setHealthObserver(healthObserver)
523+
524+
await expect(manager.ensureMounted([
525+
{
526+
provider: 'slack',
527+
mountPaths: ['/slack/channels/C123/messages']
528+
}
529+
])).rejects.toThrow('account-workspace-required')
530+
531+
expect(mock.mountInputs).toHaveLength(0)
532+
expect(healthObserver).toHaveBeenCalledWith({
533+
type: 'auth-required',
534+
reason: 'account-workspace-required',
535+
message: 'account-workspace-required'
536+
})
537+
})
538+
539+
it('does not depend on missing state files to report startup auth recovery', async () => {
540+
mock.getAccountWorkspaceId.mockRejectedValueOnce(new Error('account-workspace-required'))
541+
mock.readFile.mockRejectedValue(new Error('ENOENT'))
542+
const manager = new IntegrationMountManager()
543+
const healthObserver = vi.fn()
544+
manager.setHealthObserver(healthObserver)
545+
546+
await expect(manager.ensureMounted([
547+
{
548+
provider: 'slack',
549+
mountPaths: ['/slack/channels/C123/messages']
550+
}
551+
])).rejects.toThrow('account-workspace-required')
552+
553+
expect(mock.readFile).not.toHaveBeenCalled()
554+
expect(healthObserver).toHaveBeenCalledWith({
555+
type: 'auth-required',
556+
reason: 'account-workspace-required',
557+
message: 'account-workspace-required'
558+
})
559+
})
495560
})

0 commit comments

Comments
 (0)