1 files changed, 220 insertions, 0 deletions
diff --git a/remote/test/puppeteer/packages/puppeteer-core/src/common/NetworkEventManager.ts b/remote/test/puppeteer/packages/puppeteer-core/src/common/NetworkEventManager.ts
new file mode 100644
index 0000000000..da550a0068
--- /dev/null
+++ b/remote/test/puppeteer/packages/puppeteer-core/src/common/NetworkEventManager.ts
@@ -0,0 +1,220 @@
+/**
+ * Copyright 2022 Google Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import {Protocol} from 'devtools-protocol';
+
+import {HTTPRequest} from './HTTPRequest.js';
+
+/**
+ * @internal
+ */
+export type QueuedEventGroup = {
+  responseReceivedEvent: Protocol.Network.ResponseReceivedEvent;
+  loadingFinishedEvent?: Protocol.Network.LoadingFinishedEvent;
+  loadingFailedEvent?: Protocol.Network.LoadingFailedEvent;
+};
+
+/**
+ * @internal
+ */
+export type FetchRequestId = string;
+
+/**
+ * @internal
+ */
+export type RedirectInfo = {
+  event: Protocol.Network.RequestWillBeSentEvent;
+  fetchRequestId?: FetchRequestId;
+};
+type RedirectInfoList = RedirectInfo[];
+
+/**
+ * @internal
+ */
+export type NetworkRequestId = string;
+
+/**
+ * Helper class to track network events by request ID
+ *
+ * @internal
+ */
+export class NetworkEventManager {
+  /**
+   * There are four possible orders of events:
+   * A. `_onRequestWillBeSent`
+   * B. `_onRequestWillBeSent`, `_onRequestPaused`
+   * C. `_onRequestPaused`, `_onRequestWillBeSent`
+   * D. `_onRequestPaused`, `_onRequestWillBeSent`, `_onRequestPaused`,
+   * `_onRequestWillBeSent`, `_onRequestPaused`, `_onRequestPaused`
+   * (see crbug.com/1196004)
+   *
+   * For `_onRequest` we need the event from `_onRequestWillBeSent` and
+   * optionally the `interceptionId` from `_onRequestPaused`.
+   *
+   * If request interception is disabled, call `_onRequest` once per call to
+   * `_onRequestWillBeSent`.
+   * If request interception is enabled, call `_onRequest` once per call to
+   * `_onRequestPaused` (once per `interceptionId`).
+   *
+   * Events are stored to allow for subsequent events to call `_onRequest`.
+   *
+   * Note that (chains of) redirect requests have the same `requestId` (!) as
+   * the original request. We have to anticipate series of events like these:
+   * A. `_onRequestWillBeSent`,
+   * `_onRequestWillBeSent`, ...
+   * B. `_onRequestWillBeSent`, `_onRequestPaused`,
+   * `_onRequestWillBeSent`, `_onRequestPaused`, ...
+   * C. `_onRequestWillBeSent`, `_onRequestPaused`,
+   * `_onRequestPaused`, `_onRequestWillBeSent`, ...
+   * D. `_onRequestPaused`, `_onRequestWillBeSent`,
+   * `_onRequestPaused`, `_onRequestWillBeSent`, `_onRequestPaused`,
+   * `_onRequestWillBeSent`, `_onRequestPaused`, `_onRequestPaused`, ...
+   * (see crbug.com/1196004)
+   */
+  #requestWillBeSentMap = new Map<
+    NetworkRequestId,
+    Protocol.Network.RequestWillBeSentEvent
+  >();
+  #requestPausedMap = new Map<
+    NetworkRequestId,
+    Protocol.Fetch.RequestPausedEvent
+  >();
+  #httpRequestsMap = new Map<NetworkRequestId, HTTPRequest>();
+
+  /*
+   * The below maps are used to reconcile Network.responseReceivedExtraInfo
+   * events with their corresponding request. Each response and redirect
+   * response gets an ExtraInfo event, and we don't know which will come first.
+   * This means that we have to store a Response or an ExtraInfo for each
+   * response, and emit the event when we get both of them. In addition, to
+   * handle redirects, we have to make them Arrays to represent the chain of
+   * events.
+   */
+  #responseReceivedExtraInfoMap = new Map<
+    NetworkRequestId,
+    Protocol.Network.ResponseReceivedExtraInfoEvent[]
+  >();
+  #queuedRedirectInfoMap = new Map<NetworkRequestId, RedirectInfoList>();
+  #queuedEventGroupMap = new Map<NetworkRequestId, QueuedEventGroup>();
+
+  forget(networkRequestId: NetworkRequestId): void {
+    this.#requestWillBeSentMap.delete(networkRequestId);
+    this.#requestPausedMap.delete(networkRequestId);
+    this.#queuedEventGroupMap.delete(networkRequestId);
+    this.#queuedRedirectInfoMap.delete(networkRequestId);
+    this.#responseReceivedExtraInfoMap.delete(networkRequestId);
+  }
+
+  responseExtraInfo(
+    networkRequestId: NetworkRequestId
+  ): Protocol.Network.ResponseReceivedExtraInfoEvent[] {
+    if (!this.#responseReceivedExtraInfoMap.has(networkRequestId)) {
+      this.#responseReceivedExtraInfoMap.set(networkRequestId, []);
+    }
+    return this.#responseReceivedExtraInfoMap.get(
+      networkRequestId
+    ) as Protocol.Network.ResponseReceivedExtraInfoEvent[];
+  }
+
+  private queuedRedirectInfo(fetchRequestId: FetchRequestId): RedirectInfoList {
+    if (!this.#queuedRedirectInfoMap.has(fetchRequestId)) {
+      this.#queuedRedirectInfoMap.set(fetchRequestId, []);
+    }
+    return this.#queuedRedirectInfoMap.get(fetchRequestId) as RedirectInfoList;
+  }
+
+  queueRedirectInfo(
+    fetchRequestId: FetchRequestId,
+    redirectInfo: RedirectInfo
+  ): void {
+    this.queuedRedirectInfo(fetchRequestId).push(redirectInfo);
+  }
+
+  takeQueuedRedirectInfo(
+    fetchRequestId: FetchRequestId
+  ): RedirectInfo | undefined {
+    return this.queuedRedirectInfo(fetchRequestId).shift();
+  }
+
+  numRequestsInProgress(): number {
+    return [...this.#httpRequestsMap].filter(([, request]) => {
+      return !request.response();
+    }).length;
+  }
+
+  storeRequestWillBeSent(
+    networkRequestId: NetworkRequestId,
+    event: Protocol.Network.RequestWillBeSentEvent
+  ): void {
+    this.#requestWillBeSentMap.set(networkRequestId, event);
+  }
+
+  getRequestWillBeSent(
+    networkRequestId: NetworkRequestId
+  ): Protocol.Network.RequestWillBeSentEvent | undefined {
+    return this.#requestWillBeSentMap.get(networkRequestId);
+  }
+
+  forgetRequestWillBeSent(networkRequestId: NetworkRequestId): void {
+    this.#requestWillBeSentMap.delete(networkRequestId);
+  }
+
+  getRequestPaused(
+    networkRequestId: NetworkRequestId
+  ): Protocol.Fetch.RequestPausedEvent | undefined {
+    return this.#requestPausedMap.get(networkRequestId);
+  }
+
+  forgetRequestPaused(networkRequestId: NetworkRequestId): void {
+    this.#requestPausedMap.delete(networkRequestId);
+  }
+
+  storeRequestPaused(
+    networkRequestId: NetworkRequestId,
+    event: Protocol.Fetch.RequestPausedEvent
+  ): void {
+    this.#requestPausedMap.set(networkRequestId, event);
+  }
+
+  getRequest(networkRequestId: NetworkRequestId): HTTPRequest | undefined {
+    return this.#httpRequestsMap.get(networkRequestId);
+  }
+
+  storeRequest(networkRequestId: NetworkRequestId, request: HTTPRequest): void {
+    this.#httpRequestsMap.set(networkRequestId, request);
+  }
+
+  forgetRequest(networkRequestId: NetworkRequestId): void {
+    this.#httpRequestsMap.delete(networkRequestId);
+  }
+
+  getQueuedEventGroup(
+    networkRequestId: NetworkRequestId
+  ): QueuedEventGroup | undefined {
+    return this.#queuedEventGroupMap.get(networkRequestId);
+  }
+
+  queueEventGroup(
+    networkRequestId: NetworkRequestId,
+    event: QueuedEventGroup
+  ): void {
+    this.#queuedEventGroupMap.set(networkRequestId, event);
+  }
+
+  forgetQueuedEventGroup(networkRequestId: NetworkRequestId): void {
+    this.#queuedEventGroupMap.delete(networkRequestId);
+  }
+}