From 57996bc321e9efdf0953c378d9b18d9508e766be Mon Sep 17 00:00:00 2001
From: Frederik Bolding <frederik.bolding@gmail.com>
Date: Thu, 5 Oct 2023 14:46:35 +0200
Subject: [PATCH] Notify when using response async functions

---
 .../common/BaseSnapExecutor.test.browser.ts   | 10 +++
 .../src/common/endowments/network.test.ts     |  2 +-
 .../src/common/endowments/network.ts          | 82 ++++++++++++++++---
 3 files changed, 83 insertions(+), 11 deletions(-)

diff --git a/packages/snaps-execution-environments/src/common/BaseSnapExecutor.test.browser.ts b/packages/snaps-execution-environments/src/common/BaseSnapExecutor.test.browser.ts
index 6e372ad710..ce6f22853c 100644
--- a/packages/snaps-execution-environments/src/common/BaseSnapExecutor.test.browser.ts
+++ b/packages/snaps-execution-environments/src/common/BaseSnapExecutor.test.browser.ts
@@ -1058,6 +1058,16 @@ describe('BaseSnapExecutor', () => {
       method: 'OutboundResponse',
     });
 
+    expect(await executor.readCommand()).toStrictEqual({
+      jsonrpc: '2.0',
+      method: 'OutboundRequest',
+    });
+
+    expect(await executor.readCommand()).toStrictEqual({
+      jsonrpc: '2.0',
+      method: 'OutboundResponse',
+    });
+
     expect(await executor.readCommand()).toStrictEqual({
       id: 2,
       jsonrpc: '2.0',
diff --git a/packages/snaps-execution-environments/src/common/endowments/network.test.ts b/packages/snaps-execution-environments/src/common/endowments/network.test.ts
index addf4acda5..05bca04ef6 100644
--- a/packages/snaps-execution-environments/src/common/endowments/network.test.ts
+++ b/packages/snaps-execution-environments/src/common/endowments/network.test.ts
@@ -39,7 +39,7 @@ describe('Network endowments', () => {
       expect(factoryOptions.notify).toHaveBeenCalledWith({
         method: 'OutboundResponse',
       });
-      expect(factoryOptions.notify).toHaveBeenCalledTimes(2);
+      expect(factoryOptions.notify).toHaveBeenCalledTimes(4);
     });
 
     it('can use AbortController normally', async () => {
diff --git a/packages/snaps-execution-environments/src/common/endowments/network.ts b/packages/snaps-execution-environments/src/common/endowments/network.ts
index ba3bd48499..5f8bd952b7 100644
--- a/packages/snaps-execution-environments/src/common/endowments/network.ts
+++ b/packages/snaps-execution-environments/src/common/endowments/network.ts
@@ -12,9 +12,20 @@ class ResponseWrapper implements Response {
 
   #ogResponse: Response;
 
-  constructor(ogResponse: Response, teardownRef: { lastTeardown: number }) {
+  #onStart: () => void;
+
+  #onFinish: () => void;
+
+  constructor(
+    ogResponse: Response,
+    teardownRef: { lastTeardown: number },
+    onStart: () => void,
+    onFinish: () => void,
+  ) {
     this.#ogResponse = ogResponse;
     this.#teardownRef = teardownRef;
+    this.#onStart = onStart;
+    this.#onFinish = onFinish;
   }
 
   get body(): ReadableStream<Uint8Array> | null {
@@ -54,31 +65,64 @@ class ResponseWrapper implements Response {
   }
 
   async text() {
-    return withTeardown<string>(this.#ogResponse.text(), this as any);
+    this.#onStart();
+    try {
+      return await withTeardown<string>(this.#ogResponse.text(), this as any);
+    } finally {
+      this.#onFinish();
+    }
   }
 
   async arrayBuffer(): Promise<ArrayBuffer> {
-    return withTeardown<ArrayBuffer>(
-      this.#ogResponse.arrayBuffer(),
-      this as any,
-    );
+    this.#onStart();
+    try {
+      return await withTeardown<ArrayBuffer>(
+        this.#ogResponse.arrayBuffer(),
+        this as any,
+      );
+    } finally {
+      this.#onFinish();
+    }
   }
 
   async blob(): Promise<Blob> {
-    return withTeardown<Blob>(this.#ogResponse.blob(), this as any);
+    this.#onStart();
+    try {
+      return await withTeardown<Blob>(this.#ogResponse.blob(), this as any);
+    } finally {
+      this.#onFinish();
+    }
   }
 
   clone(): Response {
     const newResponse = this.#ogResponse.clone();
-    return new ResponseWrapper(newResponse, this.#teardownRef);
+    return new ResponseWrapper(
+      newResponse,
+      this.#teardownRef,
+      this.#onStart,
+      this.#onFinish,
+    );
   }
 
   async formData(): Promise<FormData> {
-    return withTeardown<FormData>(this.#ogResponse.formData(), this as any);
+    this.#onStart();
+    try {
+      return await withTeardown<FormData>(
+        this.#ogResponse.formData(),
+        this as any,
+      );
+    } finally {
+      this.#onFinish();
+    }
   }
 
   async json(): Promise<any> {
-    return withTeardown(this.#ogResponse.json(), this as any);
+    this.#onStart();
+    try {
+      return await withTeardown(this.#ogResponse.json(), this as any);
+    } finally {
+      this.#onFinish();
+    }
   }
 }
 
@@ -127,6 +171,22 @@ const createNetwork = ({ notify }: EndowmentFactoryOptions = {}) => {
       );
     }
 
+    let started = false;
+    const onStart = () => {
+      if (!started) {
+        started = true;
+        notify({ method: 'OutboundRequest' });
+      }
+    };
+
+    let finished = false;
+    const onFinish = () => {
+      if (!finished) {
+        finished = true;
+        notify({ method: 'OutboundResponse' });
+      }
+    };
+
     let res: Response;
     let openFetchConnection: { cancel: () => Promise<void> } | undefined;
     try {
@@ -151,6 +211,8 @@ const createNetwork = ({ notify }: EndowmentFactoryOptions = {}) => {
       res = new ResponseWrapper(
         await withTeardown(fetchPromise, teardownRef),
         teardownRef,
+        onStart,
+        onFinish,
       );
     } finally {
       if (openFetchConnection !== undefined) {