From 4bd7e6cd99b97bceef9df3ecc3704cb508759b25 Mon Sep 17 00:00:00 2001 From: hugo Date: Sun, 1 Mar 2026 13:20:16 +0100 Subject: [PATCH] fix: rollback in-memory key on store failure, add corrupted-data and rollback tests --- src/main/ipc/chatHandlers.ts | 10 +++++++-- tests/engine/ChatEngine.test.ts | 29 +++++++++++++++++++------- tests/engine/SecureKeyStore.test.ts | 22 ++++++++++++++++++- tests/ipc/chatHandlersKeychain.test.ts | 20 +++++++++++++++++- 4 files changed, 70 insertions(+), 11 deletions(-) diff --git a/src/main/ipc/chatHandlers.ts b/src/main/ipc/chatHandlers.ts index 26e9c22..6bf06fa 100644 --- a/src/main/ipc/chatHandlers.ts +++ b/src/main/ipc/chatHandlers.ts @@ -127,10 +127,16 @@ export function registerChatHandlers(): void { ipcMain.handle('chat:setApiKey', async (_, apiKey: string) => { try { const manager = await getOpenCodeManager(); + const previousKey = manager.getApiKey(); manager.setApiKey(apiKey); - // Persist to encrypted storage - await getSecureKeyStore().store('opencode_api_key', apiKey); + // Persist to encrypted storage — roll back in-memory key on failure + try { + await getSecureKeyStore().store('opencode_api_key', apiKey); + } catch (storeError) { + manager.setApiKey(previousKey); + throw storeError; + } return { success: true }; } catch (error) { diff --git a/tests/engine/ChatEngine.test.ts b/tests/engine/ChatEngine.test.ts index 3e0ba73..b10098e 100644 --- a/tests/engine/ChatEngine.test.ts +++ b/tests/engine/ChatEngine.test.ts @@ -778,19 +778,34 @@ describe('ChatEngine', () => { describe('deleteSetting', () => { it('should delete a setting by key', async () => { + let capturedTable: any; let capturedPredicate: any; - vi.mocked(mockLocalDb.delete).mockImplementation(() => ({ - where: vi.fn((predicate) => { - capturedPredicate = predicate; - return Promise.resolve(); - }), - } as any)); + vi.mocked(mockLocalDb.delete).mockImplementation((table: any) => { + capturedTable = table; + return { + where: vi.fn((predicate) => { + capturedPredicate = predicate; + return Promise.resolve(); + }), + } as any; + }); await chatEngine.deleteSetting('opencode_api_key'); - expect(mockLocalDb.delete).toHaveBeenCalled(); + expect(mockLocalDb.delete).toHaveBeenCalledTimes(1); + // Verify the correct table was targeted + expect(capturedTable).toBeDefined(); + // Verify a where predicate was passed expect(capturedPredicate).toBeDefined(); }); + + it('should not throw for nonexistent keys', async () => { + vi.mocked(mockLocalDb.delete).mockImplementation(() => ({ + where: vi.fn(() => Promise.resolve()), + } as any)); + + await expect(chatEngine.deleteSetting('nonexistent_key')).resolves.not.toThrow(); + }); }); describe('setSelectedModel', () => { diff --git a/tests/engine/SecureKeyStore.test.ts b/tests/engine/SecureKeyStore.test.ts index 42a6902..bd354a5 100644 --- a/tests/engine/SecureKeyStore.test.ts +++ b/tests/engine/SecureKeyStore.test.ts @@ -5,7 +5,7 @@ * and ChatEngine dependencies. */ -import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest'; +import { describe, it, expect, beforeEach, vi } from 'vitest'; // Track mock state let safeStorageAvailable = true; @@ -205,4 +205,24 @@ describe('SecureKeyStore', () => { expect(mockChatEngine.deleteSetting).toHaveBeenCalledWith('opencode_api_key'); }); }); + + describe('retrieve with corrupted data', () => { + it('throws when stored base64 decodes to invalid ciphertext', async () => { + const store = new SecureKeyStore(mockChatEngine as any); + + // Simulate corrupted data: valid base64 but not a valid encrypted buffer + mockSettings.set('__encrypted_api_key', Buffer.from('CORRUPT:garbage').toString('base64')); + + await expect(store.retrieve('api_key')).rejects.toThrow('Failed to decrypt'); + }); + + it('throws when stored value is not valid base64', async () => { + const store = new SecureKeyStore(mockChatEngine as any); + + // Not valid base64 — Buffer.from tolerates this but decryptString will reject it + mockSettings.set('__encrypted_api_key', '!!!not-base64!!!'); + + await expect(store.retrieve('api_key')).rejects.toThrow('Failed to decrypt'); + }); + }); }); diff --git a/tests/ipc/chatHandlersKeychain.test.ts b/tests/ipc/chatHandlersKeychain.test.ts index 1d729cf..8f7be56 100644 --- a/tests/ipc/chatHandlersKeychain.test.ts +++ b/tests/ipc/chatHandlersKeychain.test.ts @@ -258,7 +258,7 @@ describe('chatHandlers keychain integration', () => { expect(manager.setApiKey).toHaveBeenCalledWith('encrypted-stored-key'); }); - it('returns error when store() throws on chat:setApiKey', async () => { + it('returns error and rolls back in-memory key when store() throws on chat:setApiKey', async () => { secureKeyStoreStoreError = new Error('encryption unavailable'); const mod = await import('../../src/main/ipc/chatHandlers'); @@ -266,10 +266,28 @@ describe('chatHandlers keychain integration', () => { mod.initializeChatHandlers(() => mainWindowMock as never, mockBundle as any); mod.registerChatHandlers(); + // Trigger init to load the existing key + const checkHandler = registeredHandlers.get('chat:checkReady'); + await checkHandler!(undefined); + + const manager = openCodeManagerInstances[0]; + // After init, the manager has the key from SecureKeyStore + expect(manager.setApiKey).toHaveBeenCalledWith('encrypted-stored-key'); + manager.setApiKey.mockClear(); + + // getApiKey returns the current in-memory key (to be restored on rollback) + manager.getApiKey.mockReturnValue('encrypted-stored-key'); + const handler = registeredHandlers.get('chat:setApiKey'); const result = await handler!(undefined, 'sk-new-key'); expect(result.success).toBe(false); expect(result.error).toContain('encryption unavailable'); + + // setApiKey should have been called twice: + // 1) with the new key (optimistic), 2) with the old key (rollback) + expect(manager.setApiKey).toHaveBeenCalledTimes(2); + expect(manager.setApiKey).toHaveBeenNthCalledWith(1, 'sk-new-key'); + expect(manager.setApiKey).toHaveBeenNthCalledWith(2, 'encrypted-stored-key'); }); });