From 3faf870ecdb29eca665047aa1016a6d2a7c7963e Mon Sep 17 00:00:00 2001
From: matt-lewandowski <matt.lewandowski@safetyculture.io>
Date: Mon, 6 Nov 2023 14:14:30 +1000
Subject: [PATCH] IRES-508: Fix streaming interceptor (#170)

* IRES-508: Fix streaming interceptor

* update version
---
 public/connect-web-interceptor.js | 114 ++++++++++++++++--------------
 public/manifest.json              |   2 +-
 2 files changed, 63 insertions(+), 53 deletions(-)

diff --git a/public/connect-web-interceptor.js b/public/connect-web-interceptor.js
index e264df9..ac0bd58 100644
--- a/public/connect-web-interceptor.js
+++ b/public/connect-web-interceptor.js
@@ -1,57 +1,67 @@
-// This interceptor will be passed every request and response. We will take that request and response
-// and post a message to the window. This will allow us to access this message in the content script. This
-// is all to make the manifest v3 happy.
-const interceptor =
-  (next) =>
-    async (req) => {
-      return await next(req).then((resp) => {
-        if (!resp.stream) {
-          window.postMessage({
-            type: "__GRPCWEB_DEVTOOLS__",
-            methodType: "unary",
-            method: req.method.name,
-            request: req.message.toJson(),
-            response: resp.message.toJson(),
-          }, "*")
-          return resp;
-        } else {
-          return {
-            ...resp,
-            async read() {
-              const streamResp = await resp.read();
-              // "value" could be undefined when "done" is true
-              if (streamResp.value) {
-                window.postMessage({
-                  type: "__GRPCWEB_DEVTOOLS__",
-                  methodType: "server_streaming",
-                  method: req.method.name,
-                  request: req.message.toJson(),
-                  response: streamResp.value.toJson(),
-                }, "*");
-              }
-              return streamResp;
-            }
-          }
-        }
-      }).catch((e) => {
-        window.postMessage({
-          type: "__GRPCWEB_DEVTOOLS__",
-          methodType: req.stream ? "server_streaming" : "unary",
-          method: req.method.name,
-          request: req.message.toJson(),
-          response: undefined,
-          error: {
-            message: e.message,
-            code: e.code,
-          }
-        }, "*")
-        throw e;
-      });
-    };
+/**
+ * Reads the message from the stream and posts it to the window.
+ * This is a generator function that will be passed to the response stream.
+ */
+async function* readMessage(req, stream) {
+  for await (const m of stream) {
+    if (m) {
+      const resp = m.toJson?.();
+      window.postMessage({
+        type: "__GRPCWEB_DEVTOOLS__",
+        methodType: "server_streaming",
+        method: req.method.name,
+        request: req.message.toJson?.(),
+        response: resp,
+      }, "*");
+    }
+    yield m;
+  }
+}
+
+/**
+ * This interceptor will be passed every request and response. We will take that request and response
+ * and post a message to the window. This will allow us to access this message in the content script. This
+ * is all to make the manifest v3 happy.
+ */
+const interceptor = (next) => async (req) => {
+  try {
+    const resp = await next(req);
+    if (!resp.stream) {
+      window.postMessage({
+        type: "__GRPCWEB_DEVTOOLS__",
+        methodType: "unary",
+        method: req.method.name,
+        request: req.message.toJson(),
+        response: resp.message.toJson(),
+      }, "*")
+      return resp;
+    } else {
+      return {
+        ...resp,
+        message: readMessage(req, resp.message),
+      }
+    }
+  } catch (e) {
+    window.postMessage({
+      type: "__GRPCWEB_DEVTOOLS__",
+      methodType: req.stream ? "server_streaming" : "unary",
+      method: req.method.name,
+      request: req.message.toJson(),
+      response: undefined,
+      error: {
+        message: e.message,
+        code: e.code,
+      }
+    }, "*")
+    throw e;
+  }
+};
 
 window.__CONNECT_WEB_DEVTOOLS__ = interceptor;
 
-// Since we are loading inject.js as a script, the order at which it is loaded is not guaranteed.
-// So we will publish a custom event that can be used, to be used to assign the interceptor.
+/**
+ * Since we are loading inject.js as a script, the order at which it is loaded is not guaranteed.
+ * So we will publish a custom event that can be used, to be used to assign the interceptor.
+ */
 const readyEvent = new CustomEvent("connect-web-dev-tools-ready");
 window.dispatchEvent(readyEvent);
diff --git a/public/manifest.json b/public/manifest.json
index 0d98552..b402af3 100644
--- a/public/manifest.json
+++ b/public/manifest.json
@@ -1,6 +1,6 @@
 {
   "name": "gRPC-Web Developer Tools",
-  "version": "1.2.1",
+  "version": "1.4.0",
   "manifest_version": 2,
   "description": "gRPC-Web Developer Tools for debugging application's gRPC-Web network requests.",
   "icons": {