diff --git a/src/main/lib/ipc/registerTelemetryHandlers.test.ts b/src/main/lib/ipc/registerTelemetryHandlers.test.ts new file mode 100644 index 00000000..a6c5522b --- /dev/null +++ b/src/main/lib/ipc/registerTelemetryHandlers.test.ts @@ -0,0 +1,86 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ + bindUserId: vi.fn(), + capture: vi.fn(), + captureException: vi.fn(), + getFlag: vi.fn(), + handle: vi.fn(), + on: vi.fn(), + recordExposure: vi.fn(), + registerPersonProperties: vi.fn() +})) + +vi.mock('electron', () => ({ + ipcMain: { + handle: mocks.handle, + on: mocks.on + } +})) + +vi.mock('../telemetry', () => ({ + bindUserId: mocks.bindUserId, + capture: mocks.capture, + captureException: mocks.captureException, + registerPersonProperties: mocks.registerPersonProperties +})) + +vi.mock('../experiments', () => ({ + getFlag: mocks.getFlag, + recordExposure: mocks.recordExposure +})) + +import { registerTelemetryHandlers } from './registerTelemetryHandlers' + +type IpcListener = (_event: unknown, payload: unknown) => void + +function listener(channel: string): IpcListener { + const call = mocks.on.mock.calls.find(([name]) => name === channel) + expect(call).toBeDefined() + return call![1] as IpcListener +} + +describe('registerTelemetryHandlers', () => { + beforeEach(() => { + vi.clearAllMocks() + registerTelemetryHandlers() + }) + + it('caps event property keys, arrays, and strings', () => { + const properties: Record = { + long: 'x'.repeat(3000), + model_paths: Array.from({ length: 200 }, (_, i) => (i === 0 ? 'y'.repeat(3000) : i)) + } + for (let i = 0; i < 200; i++) properties[`key_${i}`] = i + + listener('telemetry:capture')(null, { event: 'comfy.desktop.test', properties }) + + const sent = mocks.capture.mock.calls[0]![1] as Record + expect(Object.keys(sent)).toHaveLength(128) + expect(sent.long).toBe('x'.repeat(2048)) + expect(sent.key_125).toBe(125) + expect(sent.key_126).toBeUndefined() + + const paths = sent.model_paths as unknown[] + expect(paths).toHaveLength(128) + expect(paths[0]).toBe('y'.repeat(2048)) + expect(paths[127]).toBe(127) + }) + + it('keeps person properties scalar-only while applying the same caps', () => { + const properties: Record = { + array: [1, 2, 3], + long: 'x'.repeat(3000) + } + for (let i = 0; i < 200; i++) properties[`key_${i}`] = i + + listener('telemetry:registerProperties')(null, properties) + + const sent = mocks.registerPersonProperties.mock.calls[0]![0] as Record + expect(Object.keys(sent)).toHaveLength(127) + expect(sent.array).toBeUndefined() + expect(sent.long).toBe('x'.repeat(2048)) + expect(sent.key_125).toBe(125) + expect(sent.key_126).toBeUndefined() + }) +}) diff --git a/src/main/lib/ipc/registerTelemetryHandlers.ts b/src/main/lib/ipc/registerTelemetryHandlers.ts index a142ddd1..fd607829 100644 --- a/src/main/lib/ipc/registerTelemetryHandlers.ts +++ b/src/main/lib/ipc/registerTelemetryHandlers.ts @@ -34,21 +34,58 @@ function isTelemetryValue(v: unknown): v is mainTelemetry.TelemetryValue { ) } -// Per-key filter to the TelemetryValue contract; renderer payloads cross a -// trust boundary, so non-primitives (incl. arrays) are dropped per-key. -function asProps(value: unknown): Record { +const MAX_TELEMETRY_KEYS = 128 +const MAX_TELEMETRY_ARRAY_ITEMS = 128 +const MAX_TELEMETRY_STRING_LENGTH = 2048 + +function clampTelemetryValue(v: mainTelemetry.TelemetryValue): mainTelemetry.TelemetryValue { + return typeof v === 'string' ? v.slice(0, MAX_TELEMETRY_STRING_LENGTH) : v +} + +function asTelemetryValueArray(v: unknown): mainTelemetry.TelemetryValue[] | null { + if (!Array.isArray(v)) return null + const out: mainTelemetry.TelemetryValue[] = [] + for (let i = 0; i < v.length && i < MAX_TELEMETRY_ARRAY_ITEMS; i++) { + const raw = v[i] + if (!isTelemetryValue(raw)) return null + out.push(clampTelemetryValue(raw)) + } + return out +} + +function asTelemetryObject(value: unknown, allowArrays: true): mainTelemetry.TelemetryContext +function asTelemetryObject( + value: unknown, + allowArrays: false +): Record +function asTelemetryObject(value: unknown, allowArrays: boolean): mainTelemetry.TelemetryContext { if (!value || typeof value !== 'object' || Array.isArray(value)) return {} - const out: Record = {} - for (const [key, raw] of Object.entries(value)) { - if (typeof key !== 'string') continue + const source = value as Record + const out: mainTelemetry.TelemetryContext = {} + let count = 0 + for (const key in source) { + if (!Object.prototype.hasOwnProperty.call(source, key)) continue + if (count++ >= MAX_TELEMETRY_KEYS) break + const raw = source[key] if (isTelemetryValue(raw)) { - out[key] = raw + out[key] = clampTelemetryValue(raw) + } else if (allowArrays) { + const array = asTelemetryValueArray(raw) + if (array) out[key] = array } - // Drop anything else (objects, arrays, functions, symbols, etc.) silently. } return out } +// Per-key filter to the TelemetryContext contract. +function asProps(value: unknown): mainTelemetry.TelemetryContext { + return asTelemetryObject(value, true) +} + +function asPersonProps(value: unknown): Record { + return asTelemetryObject(value, false) +} + export function registerTelemetryHandlers(): void { ipcMain.on('telemetry:capture', (_event, payload: CapturePayload) => { const eventName = asString(payload?.event) @@ -65,7 +102,7 @@ export function registerTelemetryHandlers(): void { }) ipcMain.on('telemetry:registerProperties', (_event, properties: unknown) => { - const props = asProps(properties) + const props = asPersonProps(properties) if (Object.keys(props).length === 0) return mainTelemetry.registerPersonProperties(props) }) @@ -76,7 +113,7 @@ export function registerTelemetryHandlers(): void { if (!payload || typeof payload !== 'object') return const userId = asString((payload as Record).userId) if (!userId) return - const properties = asProps((payload as Record).properties) + const properties = asPersonProps((payload as Record).properties) mainTelemetry.bindUserId(userId, properties) }) diff --git a/src/main/lib/telemetry.test.ts b/src/main/lib/telemetry.test.ts index 0a24af80..280f7b2c 100644 --- a/src/main/lib/telemetry.test.ts +++ b/src/main/lib/telemetry.test.ts @@ -366,6 +366,18 @@ describe('telemetry SDK-level privacy safety nets', () => { expect(msg).not.toContain('64911') }) + it('scrubs string array entries as a last-resort safety net', () => { + captured.length = 0 + telemetry.capture('comfy.desktop.execution.error', { + model_paths: ['C:\\Users\\64911\\Documents\\model.safetensors', 'LoadImage'] + }) + expect(captured).toHaveLength(1) + const paths = captured[0]!.properties?.model_paths as string[] + expect(paths[0]).toContain('[REDACTED]') + expect(paths[0]).not.toContain('64911') + expect(paths[1]).toBe('LoadImage') + }) + it('leaves non-string property types untouched', () => { captured.length = 0 telemetry.capture('comfy.desktop.execution.completed', { diff --git a/src/main/lib/telemetry.ts b/src/main/lib/telemetry.ts index d1e6dcd2..88502032 100644 --- a/src/main/lib/telemetry.ts +++ b/src/main/lib/telemetry.ts @@ -728,11 +728,17 @@ function scrubProperties(properties: TelemetryContext): TelemetryContext { let mutated: TelemetryContext | null = null for (const key of Object.keys(properties)) { const value = properties[key] - if (typeof value !== 'string') continue - const cleaned = scrubAll(value) - if (cleaned === value) continue - if (!mutated) mutated = { ...properties } - mutated[key] = cleaned + if (typeof value === 'string') { + const cleaned = scrubAll(value) + if (cleaned === value) continue + if (!mutated) mutated = { ...properties } + mutated[key] = cleaned + } else if (Array.isArray(value)) { + const cleaned = value.map((entry) => (typeof entry === 'string' ? scrubAll(entry) : entry)) + if (cleaned.every((entry, index) => entry === value[index])) continue + if (!mutated) mutated = { ...properties } + mutated[key] = cleaned + } } return mutated ?? properties } diff --git a/src/preload/comfyPreload.ts b/src/preload/comfyPreload.ts index b4ed9cf6..1e078afb 100644 --- a/src/preload/comfyPreload.ts +++ b/src/preload/comfyPreload.ts @@ -3,6 +3,7 @@ import type { IpcRendererEvent } from 'electron' import type { ComfyDesktop2Bridge, ComfyDesktop2LogsBridge, + ComfyDesktop2TelemetryBridge, ComfyDesktop2TerminalBridge, ComfyDownloadProgress, LogsOutputMsg, @@ -38,8 +39,7 @@ const Terminal: ComfyDesktop2TerminalBridge = { * the inline injection doesn't need to know its own ID. */ openPopout: (): Promise => ipcRenderer.invoke('terminal-popout-open', null), onOutput: (callback: (data: string) => void): (() => void) => { - const handler = (_event: IpcRendererEvent, payload: { data: string }) => - callback(payload.data) + const handler = (_event: IpcRendererEvent, payload: { data: string }) => callback(payload.data) ipcRenderer.on('terminal-output', handler) return () => ipcRenderer.removeListener('terminal-output', handler) }, @@ -47,7 +47,7 @@ const Terminal: ComfyDesktop2TerminalBridge = { const handler = () => callback() ipcRenderer.on('terminal-exited', handler) return () => ipcRenderer.removeListener('terminal-exited', handler) - }, + } } /** @@ -69,11 +69,27 @@ const Logs: ComfyDesktop2LogsBridge = { * so the inline injection doesn't need to know its own ID. */ openPopout: (): Promise => ipcRenderer.invoke('logs-popout-open', null), onOutput: (callback: (msg: LogsOutputMsg) => void): (() => void) => { - const handler = (_event: IpcRendererEvent, payload: LogsOutputMsg) => - callback(payload) + const handler = (_event: IpcRendererEvent, payload: LogsOutputMsg) => callback(payload) ipcRenderer.on('logs-output', handler) return () => ipcRenderer.removeListener('logs-output', handler) - }, + } +} + +const Telemetry: ComfyDesktop2TelemetryBridge = { + capture: (event, properties): void => { + // Telemetry must never break the caller. `ipcRenderer.send` throws + // synchronously on non-structured-cloneable values (functions, DOM + // nodes, symbols, circular refs); since this surface is exposed to + // the hosted ComfyUI frontend, a stray bad payload would otherwise + // propagate into unrelated frontend code. Mirrors the same + // try/catch contract as `window.api.captureTelemetry` in + // `src/preload/api.ts`. + try { + ipcRenderer.send('telemetry:capture', { event, properties }) + } catch { + // ignore: telemetry must never break the renderer + } + } } const bridge = { @@ -82,7 +98,11 @@ const bridge = { return ipcRenderer.invoke('desktop2-download-model', { url, filename, directory }) }, downloadAsset: (url: string, filename: string, authToken?: string): Promise => { - return ipcRenderer.invoke('desktop2-download-asset', { url, filename, authToken: authToken || undefined }) + return ipcRenderer.invoke('desktop2-download-asset', { + url, + filename, + authToken: authToken || undefined + }) }, pauseDownload: (url: string): Promise => { return ipcRenderer.invoke('model-download-pause', { url }) @@ -93,9 +113,7 @@ const bridge = { cancelDownload: (url: string): Promise => { return ipcRenderer.invoke('model-download-cancel', { url }) }, - onDownloadProgress: ( - callback: (data: ComfyDownloadProgress) => void - ): (() => void) => { + onDownloadProgress: (callback: (data: ComfyDownloadProgress) => void): (() => void) => { const handler = (_event: IpcRendererEvent, data: unknown) => callback(data as ComfyDownloadProgress) ipcRenderer.on('desktop2-download-progress', handler) @@ -105,7 +123,8 @@ const bridge = { ipcRenderer.send('desktop2-theme-report', { bg, text }) }, Terminal, - Logs + Logs, + Telemetry } satisfies ComfyDesktop2Bridge contextBridge.exposeInMainWorld('__comfyDesktop2', bridge)