fix(core): block resources not working

Author: blakebyrnes

core/lib/AwaitedEventListener.ts

@@ -19,6 +19,14 @@ export default class AwaitedEventListener {
     session.on('closing', () => this.triggerListenersWithType('close'));
   }
 
+  public close() {
+    for (const entry of this.listenersById.values()) {
+      if (entry.type !== 'close') {
+        this.remove(entry.id);
+      }
+    }
+  }
+
   public listen(
     sessionMeta: ISessionMeta,
     jsPath: IJsPath | null,
@@ -59,11 +67,13 @@ export default class AwaitedEventListener {
 
     const { type, meta, listenFn, jsPath } = listener;
     const tab = Session.getTab(meta);
-    if (type === 'resource') {
-      tab.off('resource', listenFn);
-    }
-    if (isWebsocketListener(listener)) {
-      tab.sessionState.stopWebsocketMessages(jsPath[1] as string, listenFn);
+    if (tab) {
+      if (type === 'resource') {
+        tab.off('resource', listenFn);
+      }
+      if (isWebsocketListener(listener)) {
+        tab.sessionState.stopWebsocketMessages(jsPath[1] as string, listenFn);
+      }
     }
 
     for (const entry of this.listenersById.values()) {

core/lib/Session.ts

@@ -215,6 +215,7 @@ export default class Session extends TypedEventEmitter<{
       sessionId: this.id,
     });
 
+    this.awaitedEventListener.close();
     await this.mitmRequestSession.close();
     await Promise.all(Object.values(this.tabs).map(x => x.close()));
     try {
@@ -289,7 +290,9 @@ export default class Session extends TypedEventEmitter<{
     }
 
     const resource = this.sessionState.captureResource(tab?.id ?? tabId, event, true);
-    tab?.emit('resource', resource);
+    if (!event.didBlockResource) {
+      tab?.emit('resource', resource);
+    }
     tab?.checkForResolvedNavigation(event.browserRequestId, resource);
   }
 

full-client/test/basic.test.ts

@@ -51,12 +51,10 @@ describe('basic Full Client tests', () => {
     await agent.activeTab.on('resource', event => resources.push(event));
     await agent.goto(`${koaServer.baseUrl}/block`);
     await agent.waitForPaintingStable();
-    expect(resources).toHaveLength(3);
-    for (const resource of resources) {
-      const status = await resource.response.statusCode;
-      if (resource.type === 'Document') expect(status).toBe(200);
-      else expect(status).toBe(404);
-    }
+    await new Promise(setImmediate);
+    expect(resources).toHaveLength(1);
+    expect(await resources[0].response.statusCode).toBe(200);
+    expect(resources[0].type).toBe('Document');
   });
 
   it('should get unreachable proxy errors in the client', async () => {

mitm/handlers/BaseHttpHandler.ts

@@ -41,20 +41,21 @@ export default abstract class BaseHttpHandler {
 
       if (session.isClosing) return context.setState(ResourceState.SessionClosed);
 
+      // need to determine resource type before blocking
+      await HeadersHandler.determineResourceType(context);
+
       if (BlockHandler.shouldBlockRequest(context)) {
         context.setState(ResourceState.Blocked);
         log.info(`Http.RequestBlocked`, {
           sessionId: session.sessionId,
           url: context.url.href,
         });
-        await HeadersHandler.determineResourceType(context);
         await context.browserHasRequested;
         session.emit('response', MitmRequestContext.toEmittedResource(this.context));
         // already wrote reply
         return;
       }
 
-      await HeadersHandler.determineResourceType(context);
       await CookieHandler.setProxyToServerCookies(context);
 
       // do one more check on the session before doing a connect

mitm/handlers/HeadersHandler.ts

@@ -179,6 +179,7 @@ export default class HeadersHandler {
   private static secFetchDestToResourceType(secFetchDest: string): ResourceType {
     switch (secFetchDest) {
       case 'document':
+      case 'iframe':
       case 'nested-document': // guess
         return 'Document';
       case 'style':