fix(mitm): don’t send duplicated headers to h2

blakebyrnes · blakebyrnes · commit ece1b1fe421d · 2020-10-08T13:58:07.000-04:00
Also catch h2 push streams closing/errors
diff --git a/full-client/test/emulate.test.ts b/full-client/test/emulate.test.ts
@@ -182,7 +182,7 @@ test('should not leave stack trace markers when calling in page functions', asyn
   const browser = await SecretAgent.createBrowser();
   koaServer.get('/marker', ctx => {
     ctx.body = `
-<body></body>
+<body>
 <script type="text/javascript">
   function errorCheck() {
     const err = new Error('This is from inside');
@@ -195,6 +195,7 @@ test('should not leave stack trace markers when calling in page functions', asyn
     };
   })(document.querySelectorAll);
 </script>
+</body>
     `;
   });
   const url = `${koaServer.baseUrl}/marker`;
diff --git a/mitm/handlers/CacheHandler.ts b/mitm/handlers/CacheHandler.ts
@@ -24,7 +24,8 @@ export default class CacheHandler {
       const cache = this.responseCache?.get(ctx.url.href);
 
       if (cache?.etag) {
-        ctx.requestHeaders['If-None-Match'] = cache.etag;
+        const key = this.ctx.isServerHttp2 ? 'if-none-match' : 'If-None-Match';
+        ctx.requestHeaders[key] = cache.etag;
         this.didProposeCachedResource = true;
       }
     }
diff --git a/mitm/handlers/HeadersHandler.ts b/mitm/handlers/HeadersHandler.ts
@@ -146,7 +146,7 @@ export default class HeadersHandler {
       }
       if (singleValueHttp2Headers.has(lowerKey)) {
         const value = ctx.requestHeaders[key];
-        if (Array.isArray(value) && value.length) ctx.requestHeaders[key] = value[0];
+        if (Array.isArray(value) && value.length) ctx.requestHeaders[lowerKey] = value[0];
       }
     }
   }
diff --git a/mitm/handlers/HttpRequestHandler.ts b/mitm/handlers/HttpRequestHandler.ts
@@ -1,7 +1,7 @@
 import * as http from 'http';
 import Log from '@secret-agent/commons/Logger';
 import * as http2 from 'http2';
-import { ClientHttp2Stream } from 'http2';
+import { ClientHttp2Stream, Http2ServerResponse } from 'http2';
 import IMitmRequestContext from '../interfaces/IMitmRequestContext';
 import HeadersHandler from './HeadersHandler';
 import CookieHandler from './CookieHandler';
@@ -144,7 +144,7 @@ export default class HttpRequestHandler extends BaseHttpHandler {
 
     for await (const chunk of serverToProxyResponse) {
       const data = context.cacheHandler.onResponseData(chunk as Buffer);
-      if (data) {
+      if (data && !(proxyToClientResponse as Http2ServerResponse).stream?.destroyed) {
         proxyToClientResponse.write(data, error => {
           if (error) this.onError('ServerToProxy.WriteResponseError', error);
         });
diff --git a/mitm/lib/MitmRequestAgent.ts b/mitm/lib/MitmRequestAgent.ts
@@ -255,6 +255,12 @@ export default class MitmRequestAgent {
     const session = this.session;
     const sessionId = session.sessionId;
     log.info('Http2Client.pushReceived', { sessionId, headers, flags });
+    stream.on('error', error => {
+      log.warn('Http2.ProxyToServer.PushStreamError', {
+        sessionId,
+        error,
+      });
+    });
 
     const pushContext = MitmRequestContext.createFromHttp2Push(parentContext, rawHeaders);
     this.session.trackResource(pushContext);
@@ -266,8 +272,17 @@ export default class MitmRequestAgent {
     }
 
     pushContext.serverToProxyResponse = stream;
-    HeadersHandler.stripHttp1HeadersForHttp2(pushContext);
 
+    // emit request
+    if (!parentContext.isClientHttp2) {
+      log.warn('Http2Client.pushReceivedWithNonH2BrowserClient', {
+        sessionId,
+        path: headers[':path'],
+      });
+      return stream.close(http2.constants.NGHTTP2_REFUSED_STREAM);
+    }
+
+    HeadersHandler.stripHttp1HeadersForHttp2(pushContext);
     const onResponseHeaders = new Promise(resolve => {
       stream.once('push', (responseHeaders, responseFlags, responseRawHeaders) => {
         MitmRequestContext.readHttp2Response(
@@ -280,60 +295,68 @@ export default class MitmRequestAgent {
       });
     });
 
-    // emit request
-    if (!parentContext.isClientHttp2) {
-      log.warn('Http2Client.pushReceivedWithNonH2BrowserClient', {
-        sessionId,
-        path: headers[':path'],
-      });
-      return stream.close(http2.constants.NGHTTP2_REFUSED_STREAM);
-    }
+    if (stream.destroyed) return;
 
     const clientToProxyRequest = parentContext.clientToProxyRequest as http2.Http2ServerRequest;
 
-    clientToProxyRequest.stream.pushStream(pushContext.requestHeaders, async (err, pushStream) => {
-      if (err) {
-        log.warn('Http2.PushStreamError', {
-          sessionId,
-          err,
+    clientToProxyRequest.stream.pushStream(
+      pushContext.requestHeaders,
+      async (error, pushStream) => {
+        if (error) {
+          log.warn('Http2.ClientToProxy.PushStreamError', {
+            sessionId,
+            error,
+          });
+          return;
+        }
+        pushStream.on('error', pushError => {
+          log.warn('Http2.ClientToProxy.PushStreamError', {
+            sessionId,
+            error: pushError,
+          });
         });
-        return;
-      }
-
-      stream.on('headers', additional => pushStream.additionalHeaders(additional));
-
-      let trailers: http2.IncomingHttpHeaders;
-      stream.once('trailers', trailerHeaders => {
-        trailers = trailerHeaders;
-      });
-
-      await onResponseHeaders;
-      pushStream.respond(pushContext.responseHeaders, { waitForTrailers: true });
-      pushStream.on('wantTrailers', async () => {
-        pushContext.responseTrailers = trailers;
-        pushStream.sendTrailers(pushContext.responseTrailers ?? {});
-      });
-      const cache = pushContext.cacheHandler;
-      cache.onHttp2PushStream();
 
-      for await (const chunk of stream) {
-        cache.onResponseData(chunk);
-        pushStream.write(chunk);
-      }
-      if (cache.shouldServeCachedData) {
-        pushStream.write(cache.cacheData);
-      }
+        stream.on('headers', additional => {
+          if (!pushStream.destroyed) pushStream.additionalHeaders(additional);
+        });
 
-      stream.end();
-      pushStream.end();
-      cache.onResponseEnd();
+        let trailers: http2.IncomingHttpHeaders;
+        stream.once('trailers', trailerHeaders => {
+          trailers = trailerHeaders;
+        });
 
-      await HeadersHandler.waitForBrowserRequest(pushContext);
-      parentContext.requestSession.emit(
-        'response',
-        MitmRequestContext.toEmittedResource(pushContext),
-      );
-    });
+        await onResponseHeaders;
+        if (pushStream.destroyed || stream.destroyed) {
+          return;
+        }
+        pushStream.respond(pushContext.responseHeaders, { waitForTrailers: true });
+        pushStream.on('wantTrailers', async () => {
+          pushContext.responseTrailers = trailers;
+          pushStream.sendTrailers(pushContext.responseTrailers ?? {});
+        });
+        const cache = pushContext.cacheHandler;
+        cache.onHttp2PushStream();
+
+        for await (const chunk of stream) {
+          if (pushStream.destroyed || stream.destroyed) return;
+          cache.onResponseData(chunk);
+          pushStream.write(chunk);
+        }
+        if (cache.shouldServeCachedData) {
+          if (!pushStream.destroyed && !stream.destroyed) pushStream.write(cache.cacheData);
+        }
+
+        if (!stream.destroyed) stream.end();
+        if (!pushStream.destroyed) pushStream.end();
+        cache.onResponseEnd();
+
+        await HeadersHandler.waitForBrowserRequest(pushContext);
+        parentContext.requestSession.emit(
+          'response',
+          MitmRequestContext.toEmittedResource(pushContext),
+        );
+      },
+    );
   }
 
   private getHttp2Session(origin: string) {
diff --git a/mitm/test/basic.test.ts b/mitm/test/basic.test.ts
@@ -1,16 +1,16 @@
-import http from "http";
-import { Helpers } from "@secret-agent/testing";
-import HttpProxyAgent from "http-proxy-agent";
-import { AddressInfo } from "net";
-import WebSocket from "ws";
-import Url from "url";
-import { createPromise } from "@secret-agent/commons/utils";
-import HttpRequestHandler from "../handlers/HttpRequestHandler";
-import RequestSession from "../handlers/RequestSession";
-import MitmServer from "../lib/MitmProxy";
-import HeadersHandler from "../handlers/HeadersHandler";
-import HttpUpgradeHandler from "../handlers/HttpUpgradeHandler";
-import { parseRawHeaders } from "../lib/Utils";
+import http from 'http';
+import { Helpers } from '@secret-agent/testing';
+import HttpProxyAgent from 'http-proxy-agent';
+import { AddressInfo } from 'net';
+import WebSocket from 'ws';
+import Url from 'url';
+import { createPromise } from '@secret-agent/commons/utils';
+import HttpRequestHandler from '../handlers/HttpRequestHandler';
+import RequestSession from '../handlers/RequestSession';
+import MitmServer from '../lib/MitmProxy';
+import HeadersHandler from '../handlers/HeadersHandler';
+import HttpUpgradeHandler from '../handlers/HttpUpgradeHandler';
+import { parseRawHeaders } from '../lib/Utils';
 
 const mocks = {
   httpRequestHandler: {
@@ -35,14 +35,16 @@ beforeEach(() => {
 afterAll(Helpers.afterAll);
 afterEach(Helpers.afterEach);
 
+let sessionCounter = 1;
+
 describe('basic MitM tests', () => {
   it('should send request through proxy', async () => {
     const httpServer = await Helpers.runHttpServer();
     const mitmServer = await MitmServer.start();
     Helpers.needsClosing.push(mitmServer);
     const proxyHost = `http://localhost:${mitmServer.port}`;
 
-    const session = new RequestSession('1', 'any agent', null);
+    const session = new RequestSession(`${(sessionCounter += 1)}`, 'any agent', null);
 
     const proxyCredentials = session.getProxyCredentials();
     expect(mocks.httpRequestHandler.onRequest).toBeCalledTimes(0);
@@ -69,7 +71,7 @@ describe('basic MitM tests', () => {
     const proxyCredentials = session.getProxyCredentials();
     expect(mocks.httpRequestHandler.onRequest).toBeCalledTimes(0);
 
-    let rawHeaders: string[];
+    let rawHeaders: string[] = null;
     const res = await Helpers.httpRequest(
       httpServer.url,
       'GET',
@@ -97,7 +99,7 @@ describe('basic MitM tests', () => {
     Helpers.needsClosing.push(mitmServer);
     const proxyHost = `http://localhost:${mitmServer.port}`;
 
-    const session = new RequestSession('1', 'any agent', null);
+    const session = new RequestSession(`${(sessionCounter += 1)}`, 'any agent', null);
 
     const proxyCredentials = session.getProxyCredentials();
     expect(mocks.httpRequestHandler.onRequest).toBeCalledTimes(0);
@@ -123,7 +125,11 @@ describe('basic MitM tests', () => {
       socket.end();
     });
 
-    const session = new RequestSession('1', 'any agent', Promise.resolve(upstreamProxyHost));
+    const session = new RequestSession(
+      `${(sessionCounter += 1)}`,
+      'any agent',
+      Promise.resolve(upstreamProxyHost),
+    );
 
     const proxyCredentials = session.getProxyCredentials();
 
@@ -156,7 +162,7 @@ describe('basic MitM tests', () => {
     Helpers.needsClosing.push(mitmServer);
     const proxyHost = `http://localhost:${mitmServer.port}`;
 
-    const session = new RequestSession('1', 'any agent', null);
+    const session = new RequestSession(`${(sessionCounter += 1)}`, 'any agent', null);
 
     const proxyCredentials = session.getProxyCredentials();
     expect(mocks.httpRequestHandler.onRequest).toBeCalledTimes(0);
@@ -185,7 +191,7 @@ describe('basic MitM tests', () => {
     Helpers.needsClosing.push(mitmServer);
     const proxyHost = `http://localhost:${mitmServer.port}`;
 
-    const session = new RequestSession('2', 'any agent', null);
+    const session = new RequestSession(`${(sessionCounter += 1)}`, 'any agent', null);
 
     const proxyCredentials = session.getProxyCredentials();
 
@@ -204,7 +210,7 @@ describe('basic MitM tests', () => {
     Helpers.needsClosing.push(mitmServer);
     const proxyHost = `http://localhost:${mitmServer.port}`;
 
-    const session = new RequestSession('3', 'any agent', null);
+    const session = new RequestSession(`${(sessionCounter += 1)}`, 'any agent', null);
 
     const proxyCredentials = session.getProxyCredentials();
 
@@ -278,7 +284,7 @@ describe('basic MitM tests', () => {
     const serverMessages = [];
     const serverMessagePromise = createPromise();
     const wsServer = new WebSocket.Server({ noServer: true });
-    const session = new RequestSession('4', 'any agent', null);
+    const session = new RequestSession(`${(sessionCounter += 1)}`, 'any agent', null);
 
     httpServer.server.on('upgrade', (request, socket, head) => {
       // ensure header is stripped
@@ -337,11 +343,11 @@ describe('basic MitM tests', () => {
     Helpers.needsClosing.push(mitmServer);
     const proxyHost = `http://localhost:${mitmServer.port}`;
 
-    const session = new RequestSession('1', 'any agent', null);
+    const session = new RequestSession(`${(sessionCounter += 1)}`, 'any agent', null);
     session.delegate.modifyHeadersBeforeSend = jest.fn();
     session.registerResource({
       tabId: '1',
-      browserRequestId: '1',
+      browserRequestId: '25.123',
       url: `${httpServer.url}page1`,
       method: 'GET',
       resourceType: 'Document',

Original file line number	Diff line number	Diff line change
`@@ -24,7 +24,8 @@ export default class CacheHandler {`
`24`	`24`	`const cache = this.responseCache?.get(ctx.url.href);`
`25`	`25`
`26`	`26`	`if (cache?.etag) {`
`27`		`- ctx.requestHeaders['If-None-Match'] = cache.etag;`
	`27`	`+ const key = this.ctx.isServerHttp2 ? 'if-none-match' : 'If-None-Match';`
	`28`	`+ ctx.requestHeaders[key] = cache.etag;`
`28`	`29`	`this.didProposeCachedResource = true;`
`29`	`30`	`}`
`30`	`31`	`}`
Original file line number	Diff line number	Diff line change
`@@ -146,7 +146,7 @@ export default class HeadersHandler {`
`146`	`146`	`}`
`147`	`147`	`if (singleValueHttp2Headers.has(lowerKey)) {`
`148`	`148`	`const value = ctx.requestHeaders[key];`
`149`		`- if (Array.isArray(value) && value.length) ctx.requestHeaders[key] = value[0];`
	`149`	`+ if (Array.isArray(value) && value.length) ctx.requestHeaders[lowerKey] = value[0];`
`150`	`150`	`}`
`151`	`151`	`}`
`152`	`152`	`}`