diff --git a/src/background/background.ts b/src/background/background.ts index 54081de..8554a7c 100644 --- a/src/background/background.ts +++ b/src/background/background.ts @@ -23,6 +23,7 @@ import { toEmail, EMAIL_ATTRIBUTE_TYPE, findHtmlBody } from "../lib/utils"; import { getOrCreateLocalFolder } from "../lib/folders"; import { isPGEncrypted, wasPGEncrypted } from "../lib/detection"; import { dispatchRuntimeMessage } from "./runtime-router"; +import { handleAfterSend } from "./sent-copy"; import type { Policy, SerializedRecipient, @@ -102,37 +103,13 @@ browser.runtime.onMessage.addListener( browser.compose.onBeforeSend.addListener(handleBeforeSend); -browser.compose.onAfterSend.addListener(async (tab, sendInfo) => { - const state = composeTabs.get(tab.id); - if (!state?.sentMimeData) return; - - try { - for (const msg of sendInfo.messages) { - if (await isPGEncrypted(msg.id)) { - const localFolder = await getOrCreateLocalFolder("PostGuard Sent"); - if (localFolder) { - const file = new File([state.sentMimeData as BlobPart], "sent.eml", { - type: "text/plain", - }); - const localMsg = await browser.messages.import( - file, - localFolder.id - ); - await browser.messages.move([localMsg.id], msg.folder.id); - await browser.messages.delete([msg.id], true); - } - } - } - } catch (e) { - console.error("[PostGuard] Failed to manage sent copy:", e); - notifyError("sentCopyError"); - } finally { - cleanupComposeTab(tab.id); - clearInFlightUpload(tab.id); - persistEncryptState().catch(console.warn); - persistInFlightUploads().catch(console.warn); - } -}); +browser.compose.onAfterSend.addListener((tab, sendInfo) => + handleAfterSend(tab, sendInfo, { + notifyError, + isPGEncrypted, + getOrCreateLocalFolder, + }), +); // Clean up decryptedMessages when messages are deleted browser.messages.onDeleted.addListener((deletedMessages) => { diff --git a/src/background/sent-copy.ts b/src/background/sent-copy.ts new file mode 100644 index 0000000..e04cc54 --- /dev/null +++ b/src/background/sent-copy.ts @@ -0,0 +1,71 @@ +/// + +import { + composeTabs, + cleanupComposeTab, + clearInFlightUpload, + persistEncryptState, + persistInFlightUploads, +} from "./state"; + +// Pulled out of background.ts so the post-send cleanup can be exercised +// under vitest without standing up the whole background script. Deps that +// hit Thunderbird APIs we cannot easily mock from the browser-mock helper +// are injected, and the rest goes through the browser-mock surface. + +export interface SentMessage { + id: number; + folder: { id: unknown }; +} + +export interface SentInfo { + messages: SentMessage[]; +} + +export interface SentCopyDeps { + notifyError: (messageKey: string) => void; + isPGEncrypted: (msgId: number) => Promise; + getOrCreateLocalFolder: ( + name: string, + ) => Promise<{ id: unknown } | undefined>; +} + +/** + * Post-send handler. For each message the user just sent, if it is a + * PostGuard ciphertext, swap the encrypted copy in Sent with the plaintext + * MIME stashed during encryption. State cleanup happens in `finally` so a + * failure mid-swap never strands the compose tab in `composeTabs`. + */ +export async function handleAfterSend( + tab: { id: number }, + sendInfo: SentInfo, + deps: SentCopyDeps, +): Promise { + const state = composeTabs.get(tab.id); + if (!state?.sentMimeData) return; + + try { + for (const msg of sendInfo.messages) { + if (!(await deps.isPGEncrypted(msg.id))) continue; + const localFolder = await deps.getOrCreateLocalFolder("PostGuard Sent"); + if (!localFolder) continue; + const file = new File([state.sentMimeData as BlobPart], "sent.eml", { + type: "text/plain", + }); + const localMsg = await browser.messages.import( + file, + localFolder.id as any, + ); + await browser.messages.move([localMsg.id], msg.folder.id as any); + await browser.messages.delete([msg.id], true); + } + } catch (e) { + console.error("[PostGuard] Failed to manage sent copy:", e); + deps.notifyError("sentCopyError"); + } finally { + cleanupComposeTab(tab.id); + clearInFlightUpload(tab.id); + persistEncryptState().catch(console.warn); + persistInFlightUploads().catch(console.warn); + } +} diff --git a/tests/sent-copy.test.ts b/tests/sent-copy.test.ts index f1b2970..82a2a33 100644 --- a/tests/sent-copy.test.ts +++ b/tests/sent-copy.test.ts @@ -1,22 +1,141 @@ -import { describe, it, expect } from "vitest"; +import { describe, it, expect, beforeEach, vi } from "vitest"; +import { handleAfterSend, type SentInfo } from "../src/background/sent-copy"; +import { + composeTabs, + decryptedMessages, + inFlightUploads, + pendingCryptoPopups, + pendingPolicyEditors, +} from "../src/background/state"; +import { installBrowserMock, type BrowserMock } from "./helpers/browser-mock"; -// These tests verify the onAfterSend sent copy management. -// They require browser API mocks. +let mock: BrowserMock; +let notifyError: ReturnType; +let isPGEncrypted: ReturnType; +let getOrCreateLocalFolder: ReturnType; +let importSpy: ReturnType; +let moveSpy: ReturnType; +let deleteSpy: ReturnType; + +const SENT_FOLDER = { id: "sent://Sent" }; +const LOCAL_FOLDER = { id: "local://PostGuard Sent" }; + +function deps() { + return { notifyError, isPGEncrypted, getOrCreateLocalFolder }; +} + +function sentInfo(messageIds: number[] = [1]): SentInfo { + return { + messages: messageIds.map((id) => ({ id, folder: SENT_FOLDER })), + }; +} + +beforeEach(() => { + mock = installBrowserMock(); + composeTabs.clear(); + decryptedMessages.clear(); + pendingCryptoPopups.clear(); + pendingPolicyEditors.clear(); + inFlightUploads.clear(); + + notifyError = vi.fn(); + isPGEncrypted = vi.fn(async () => true); + getOrCreateLocalFolder = vi.fn(async () => LOCAL_FOLDER); + + importSpy = vi.spyOn(browser.messages, "import"); + moveSpy = vi.spyOn(browser.messages, "move"); + deleteSpy = vi.spyOn(browser.messages, "delete"); +}); describe("onAfterSend — sent copy management", () => { - it.todo("should skip when no sentMimeData is stored"); + it("should skip when no sentMimeData is stored", async () => { + composeTabs.set(7, { encrypt: true }); + await handleAfterSend({ id: 7 }, sentInfo([1]), deps()); + + expect(isPGEncrypted).not.toHaveBeenCalled(); + expect(importSpy).not.toHaveBeenCalled(); + }); + + it("should skip non-PG-encrypted sent messages", async () => { + composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) }); + isPGEncrypted.mockResolvedValueOnce(false); + + await handleAfterSend({ id: 7 }, sentInfo([42]), deps()); + + expect(isPGEncrypted).toHaveBeenCalledWith(42); + expect(getOrCreateLocalFolder).not.toHaveBeenCalled(); + expect(importSpy).not.toHaveBeenCalled(); + }); + + it("should import plaintext MIME into PostGuard Sent folder", async () => { + composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) }); + + await handleAfterSend({ id: 7 }, sentInfo([42]), deps()); + + expect(getOrCreateLocalFolder).toHaveBeenCalledWith("PostGuard Sent"); + expect(importSpy).toHaveBeenCalledTimes(1); + const [file, folder] = importSpy.mock.calls[0] as [File, unknown]; + expect(file).toBeInstanceOf(File); + expect(file.name).toBe("sent.eml"); + expect(folder).toBe(LOCAL_FOLDER.id); + }); + + it("should move imported message to original sent folder", async () => { + composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) }); + + await handleAfterSend({ id: 7 }, sentInfo([42]), deps()); + + expect(moveSpy).toHaveBeenCalledWith([9999], SENT_FOLDER.id); + }); + + it("should delete the encrypted copy from sent folder", async () => { + composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) }); + + await handleAfterSend({ id: 7 }, sentInfo([42]), deps()); + + expect(deleteSpy).toHaveBeenCalledWith([42], true); + }); + + it("should notify user when sent copy management fails", async () => { + composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) }); + importSpy.mockRejectedValueOnce(new Error("import boom")); + + await handleAfterSend({ id: 7 }, sentInfo([42]), deps()); + + expect(notifyError).toHaveBeenCalledWith("sentCopyError"); + }); + + it("should clean up compose tab state after send", async () => { + composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) }); + inFlightUploads.set(7, { uuid: "u", recoveryToken: "r", startedAt: Date.now() }); + + await handleAfterSend({ id: 7 }, sentInfo([42]), deps()); - it.todo("should skip non-PG-encrypted sent messages"); + expect(composeTabs.has(7)).toBe(false); + expect(inFlightUploads.has(7)).toBe(false); + }); - it.todo("should import plaintext MIME into PostGuard Sent folder"); + it("should clean up compose tab state even on failure", async () => { + composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) }); + inFlightUploads.set(7, { uuid: "u", recoveryToken: "r", startedAt: Date.now() }); + importSpy.mockRejectedValueOnce(new Error("nope")); - it.todo("should move imported message to original sent folder"); + await handleAfterSend({ id: 7 }, sentInfo([42]), deps()); - it.todo("should delete the encrypted copy from sent folder"); + expect(composeTabs.has(7)).toBe(false); + expect(inFlightUploads.has(7)).toBe(false); + }); - it.todo("should notify user when sent copy management fails"); + it("should skip the swap when no local folder is available", async () => { + composeTabs.set(7, { encrypt: true, sentMimeData: new Uint8Array([1, 2, 3]) }); + getOrCreateLocalFolder.mockResolvedValueOnce(undefined); - it.todo("should clean up compose tab state after send"); + await handleAfterSend({ id: 7 }, sentInfo([42]), deps()); - it.todo("should clean up compose tab state even on failure"); + expect(importSpy).not.toHaveBeenCalled(); + expect(moveSpy).not.toHaveBeenCalled(); + expect(deleteSpy).not.toHaveBeenCalled(); + // Cleanup still runs. + expect(composeTabs.has(7)).toBe(false); + }); });