From 5d7500ffd7649ab40bbce669268063fe0783fa4e Mon Sep 17 00:00:00 2001
From: George Fu <kuhe@users.noreply.github.com>
Date: Mon, 16 Dec 2024 15:54:53 -0500
Subject: [PATCH] fix(xhr-http-handler): fix abort signal sharing for multiple
 requests (#6740)

---
 .../src/xhr-http-handler.spec.ts              | 72 +++++++++++++++----
 .../xhr-http-handler/src/xhr-http-handler.ts  | 13 +++-
 2 files changed, 70 insertions(+), 15 deletions(-)

diff --git a/packages/xhr-http-handler/src/xhr-http-handler.spec.ts b/packages/xhr-http-handler/src/xhr-http-handler.spec.ts
index f37f0619746d..5960824d74c8 100644
--- a/packages/xhr-http-handler/src/xhr-http-handler.spec.ts
+++ b/packages/xhr-http-handler/src/xhr-http-handler.spec.ts
@@ -1,4 +1,3 @@
-import { AbortSignal } from "@smithy/abort-controller";
 import { HttpRequest } from "@smithy/protocol-http";
 import { afterAll, afterEach, beforeAll, describe, expect, test as it, vi } from "vitest";
 
@@ -36,13 +35,19 @@ class XhrMock {
     this.captureArgs("getAllResponseHeaders")();
     return `responseHeaderKey: responseHeaderValue\r\nrk2: rv2`;
   }
-  setRequestHeader = this.captureArgs("setRequestHeader");
-  open = this.captureArgs("open");
+  setRequestHeader(...args: any[]) {
+    return this.captureArgs("setRequestHeader")(...args);
+  }
+  open(...args: any[]) {
+    return this.captureArgs("open")(...args);
+  }
   send(...args: any[]) {
     this.captureArgs("send")(...args);
     this.eventListeners["readystatechange"][0]();
   }
-  abort = this.captureArgs("abort");
+  abort(...args: any[]) {
+    return this.captureArgs("abort")(...args);
+  }
   addEventListener(...args: any[]) {
     this.captureArgs("addEventListener")(...args);
     const [event, callback] = args;
@@ -135,9 +140,9 @@ describe(XhrHttpHandler.name, () => {
 
   it("should respond to AbortSignal", async () => {
     const handler = new XhrHttpHandler();
-    const abortSignal = new AbortSignal();
+    const abortController = new AbortController();
 
-    await handler.handle(
+    const p1 = handler.handle(
       new HttpRequest({
         method: "PUT",
         hostname: "localhost",
@@ -148,21 +153,22 @@ describe(XhrHttpHandler.name, () => {
         protocol: "http:",
         path: "/api",
       }),
-      { abortSignal }
+      { abortSignal: abortController.signal }
     );
 
     try {
-      abortSignal.abort();
+      abortController.abort();
+      await p1;
     } catch (e) {
       expect(e.toString()).toContain("Request aborted");
     }
 
     expect(XhrMock.captures).toEqual([
-      ["upload.addEventListener", "progress", expect.any(Function)],
-      ["addEventListener", "progress", expect.any(Function)],
-      ["addEventListener", "error", expect.any(Function)],
-      ["addEventListener", "timeout", expect.any(Function)],
-      ["addEventListener", "readystatechange", expect.any(Function)],
+      ["upload.addEventListener", "progress", expect.anything()],
+      ["addEventListener", "progress", expect.anything()],
+      ["addEventListener", "error", expect.anything()],
+      ["addEventListener", "timeout", expect.anything()],
+      ["addEventListener", "readystatechange", expect.anything()],
       ["open", "PUT", "http://localhost:3000/api?k=v"],
       ["setRequestHeader", "h", "1"],
       ["send", "hello"],
@@ -171,6 +177,46 @@ describe(XhrHttpHandler.name, () => {
     ]);
   });
 
+  it("should allow an AbortSignal to abort multiple requests", async () => {
+    const handler = new XhrHttpHandler();
+    const abortController = new AbortController();
+
+    expect(abortController.signal.addEventListener).toBeInstanceOf(Function);
+
+    const xhrs = [] as XMLHttpRequest[];
+    handler.on(XhrHttpHandler.EVENTS.BEFORE_XHR_SEND, (xhr) => {
+      xhrs.push(xhr);
+    });
+
+    const request = () =>
+      handler.handle(
+        new HttpRequest({
+          method: "PUT",
+          hostname: "localhost",
+          port: 3000,
+          query: { k: "v" },
+          headers: { h: "1" },
+          body: "hello",
+          protocol: "http:",
+          path: "/api",
+        }),
+        { abortSignal: abortController.signal }
+      );
+
+    const p1 = request().catch((_) => _);
+    const p2 = request().catch((_) => _);
+    const p3 = request().catch((_) => _);
+    abortController.abort();
+    await p1;
+    await p2;
+    await p3;
+    await request().catch((_) => _);
+    await request().catch((_) => _);
+
+    expect(xhrs.length).toEqual(3);
+    expect(XhrMock.captures.filter(([source]) => source === "abort")).toEqual([["abort"], ["abort"], ["abort"]]);
+  });
+
   it("should ignore forbidden request headers", async () => {
     const handler = new XhrHttpHandler();
 
diff --git a/packages/xhr-http-handler/src/xhr-http-handler.ts b/packages/xhr-http-handler/src/xhr-http-handler.ts
index a5637056a4e0..9d551ce35f20 100644
--- a/packages/xhr-http-handler/src/xhr-http-handler.ts
+++ b/packages/xhr-http-handler/src/xhr-http-handler.ts
@@ -216,19 +216,28 @@ export class XhrHttpHandler extends EventEmitter implements HttpHandler<XhrHttpH
       }),
       requestTimeout(requestTimeoutInMs),
     ];
+    let removeSignalEventListener = () => {};
     if (abortSignal) {
       raceOfPromises.push(
         new Promise<never>((resolve, reject) => {
-          abortSignal.onabort = () => {
+          const onAbort = () => {
             xhr.abort();
             const abortError = new Error("Request aborted");
             abortError.name = "AbortError";
             reject(abortError);
           };
+          if (typeof (abortSignal as any).addEventListener === "function") {
+            const signal = abortSignal as any;
+            signal.addEventListener("abort", onAbort, { once: true });
+            removeSignalEventListener = () => signal.removeEventListener("abort", onAbort);
+          } else {
+            // backwards compatibility
+            abortSignal.onabort = onAbort;
+          }
         })
       );
     }
-    return Promise.race(raceOfPromises);
+    return Promise.race(raceOfPromises).finally(removeSignalEventListener);
   }
 
   /**