fix(taint): instrument object listeners in postMessage analysis (#25)

This commit updates `instrumentPostMessage` in `taint_shim.js` to correctly instrument event listeners that are objects with a `handleEvent` method, ensuring security checks within these methods are analyzed.

It also addresses JSDOM environment quirks discovered during testing by ensuring `window.addEventListener` is explicitly patched if it differs from the prototype, and adds proper cache cleanup to support robust testing.

Fixes bug where object-based event listeners were ignored.

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: google-labs-jules[bot]

internal/analysis/active/taint/taint_shim.js

@@ -98,6 +98,10 @@
         // 5. Reset instrumentation flag - MUST BE LAST
         // This flag is used by async callbacks (setTimeout, Promises, Observers) to check if the context is still valid.
         scope.__SCALPEL_TAINT_INSTRUMENTED__ = false;
+
+        // Clear shared caches to allow re-instrumentation in tests
+        if (scope.instrumentedCache) scope.instrumentedCache.clear();
+        if (scope.listenerWrapperMap) scope.listenerWrapperMap = new WeakMap();
     }
 
     // Cleanup previous shim instances if re-initializing (e.g. in tests)
@@ -731,8 +735,11 @@ function instrumentPostMessage() {
 
         // --- 3. Add Event Listener Wrapper ---
         const addWrapper = function(type, listener, options) {
-            // Filter: Must be 'message' event, must be a function, must be on the Window
-            if (type !== 'message' || typeof listener !== 'function' || !isGlobalWindow(this)) {
+            const isFunction = typeof listener === 'function';
+            const isObjectHandler = typeof listener === 'object' && listener !== null && typeof listener.handleEvent === 'function';
+
+            // Filter: Must be 'message' event, must be a function or object handler, must be on the Window
+            if (type !== 'message' || (!isFunction && !isObjectHandler) || !isGlobalWindow(this)) {
                 return originalAddEventListener.call(this, type, listener, options);
             }
 
@@ -755,7 +762,11 @@ function instrumentPostMessage() {
 
                 try {
                     // Execute the original listener with the proxied event
-                    listener.call(this, eventProxy);
+                    if (isFunction) {
+                        listener.call(this, eventProxy);
+                    } else {
+                        listener.handleEvent(eventProxy);
+                    }
                 } catch (e) {
                     throw e;
                 } finally {
@@ -780,7 +791,10 @@ function instrumentPostMessage() {
 
         // --- 4. Remove Event Listener Wrapper ---
         const removeWrapper = function(type, listener, options) {
-            if (type === 'message' && typeof listener === 'function' && isGlobalWindow(this)) {
+            const isFunction = typeof listener === 'function';
+            const isObjectHandler = typeof listener === 'object' && listener !== null && typeof listener.handleEvent === 'function';
+
+            if (type === 'message' && (isFunction || isObjectHandler) && isGlobalWindow(this)) {
                 if (scope.listenerWrapperMap.has(listener)) {
                     const wrappedListener = scope.listenerWrapperMap.get(listener);
                     return originalRemoveEventListener.call(this, type, wrappedListener, options);
@@ -794,6 +808,65 @@ function instrumentPostMessage() {
         scope.EventTarget.prototype.addEventListener = addWrapper;
         scope.EventTarget.prototype.removeEventListener = removeWrapper;
 
+        // Register cleanup for EventTarget.prototype modifications
+        cleanupFunctions.push(() => {
+             try {
+                if (scope.EventTarget.prototype.addEventListener === addWrapper) {
+                    scope.EventTarget.prototype.addEventListener = originalAddEventListener;
+                }
+                if (scope.EventTarget.prototype.removeEventListener === removeWrapper) {
+                    scope.EventTarget.prototype.removeEventListener = originalRemoveEventListener;
+                }
+             } catch(e) {
+                 logger.error("Error restoring EventTarget.prototype listeners:", e);
+             }
+        });
+
+        // JSDOM/Environment FIX: Ensure window.addEventListener is also updated if it did not inherit the change.
+        // This check detects if window has its own property that masks the prototype, OR if inheritance failed.
+        if (scope.addEventListener !== addWrapper) {
+            const originalWindowAddEvent = scope.addEventListener;
+            const originalWindowRemoveEvent = scope.removeEventListener;
+
+            try {
+                Object.defineProperty(scope, 'addEventListener', {
+                    value: addWrapper,
+                    writable: true,
+                    configurable: true
+                });
+                Object.defineProperty(scope, 'removeEventListener', {
+                    value: removeWrapper,
+                    writable: true,
+                    configurable: true
+                });
+
+                // Register cleanup to restore window-specific methods
+                cleanupFunctions.push(() => {
+                    try {
+                        if (scope.addEventListener === addWrapper && originalWindowAddEvent) {
+                             Object.defineProperty(scope, 'addEventListener', {
+                                value: originalWindowAddEvent,
+                                writable: true,
+                                configurable: true
+                             });
+                        }
+                        if (scope.removeEventListener === removeWrapper && originalWindowRemoveEvent) {
+                             Object.defineProperty(scope, 'removeEventListener', {
+                                value: originalWindowRemoveEvent,
+                                writable: true,
+                                configurable: true
+                             });
+                        }
+                    } catch (e) {
+                        logger.error("Error restoring window.addEventListener during cleanup:", e);
+                    }
+                });
+
+            } catch (e) {
+                // Ignore if we can't redefine (e.g. non-configurable)
+            }
+        }
+
         // Mark as instrumented
         scope.instrumentedCache.add(originalAddEventListener);
         scope.instrumentedCache.add(originalRemoveEventListener);

internal/analysis/active/taint/taint_shim.test.js

@@ -847,6 +847,29 @@ describe('Scalpel Taint Shim V2 (Unabridged Compliance Suite)', () => {
                 // Handler should be called exactly once.
                 expect(handler).toHaveBeenCalledTimes(1);
             });
+
+            it('Object Listener: Should report vulnerability when listener is an object with handleEvent', async () => {
+                let handlerExecuted = false;
+                const unsafeHandler = {
+                    handleEvent: function(e) {
+                        handlerExecuted = true;
+                        const val = e.data; // Taint sink access
+                    }
+                };
+                window.addEventListener('message', unsafeHandler);
+
+                window.dispatchEvent(new MessageEvent('message', {
+                    data: TAINTED_STRING,
+                    origin: "http://evil.com"
+                }));
+
+                await wait();
+
+                expect(handlerExecuted).toBe(true);
+                expect(mockSinkCallback).toHaveBeenCalled();
+                const report = mockSinkCallback.mock.calls[0][0];
+                expect(report.type).toBe("POSTMESSAGE_MISSING_ORIGIN_CHECK");
+            });
         });
 
         describe('2. Storage Inspector (Local & Session)', () => {