From 73b6943321f228454b5a54933e9fec6dab54b497 Mon Sep 17 00:00:00 2001
From: Derrick Hammer <derrick@derrickhammer.com>
Date: Mon, 23 Oct 2023 19:44:27 -0400
Subject: [PATCH] refactor: major rewrite to send framed messages in an array
 set, as well as do internal object hacks to ensure we intercept all headers

---
 src/index.ts | 133 ++++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 121 insertions(+), 12 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 6f6bfa1..ce4a09a 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -3,7 +3,8 @@ import { NodeHttpTransport } from "@improbable-eng/grpc-web-node-http-transport"
 import { BadgeGenerator } from "@lavanet/lava-sdk/bin/src/grpc_web_services/lavanet/lava/pairing/badges_pb_service.js";
 import { GenerateBadgeResponse } from "@lavanet/lava-sdk/bin/src/grpc_web_services/lavanet/lava/pairing/badges_pb.js";
 import { PluginAPI } from "@lumeweb/interface-relay";
-import ProtobufMessage = grpc.ProtobufMessage;
+import Metadata = grpc.Metadata;
+import client = grpc.client;
 
 interface BadgeRequest {
   data: Uint8Array;
@@ -46,15 +47,85 @@ function unframeRequest(frame: Uint8Array): Uint8Array | null {
   return frame.subarray(5);
 }
 
-function frameRequest(request: ProtobufMessage): Uint8Array {
-  const bytes = request.serializeBinary();
-  const frame = new ArrayBuffer(bytes.byteLength + 5);
-  new DataView(frame, 1, 4).setUint32(0, bytes.length, false /* big endian */);
-  new Uint8Array(frame, 5).set(bytes);
-  return new Uint8Array(frame);
+const HEADER_SIZE = 5;
+
+function frameRequest(data: Uint8Array, type = 0x0): Uint8Array {
+  // Serialize the request if needed.
+  // Create the frame with the necessary length.
+  const frame = new Uint8Array(data.byteLength + HEADER_SIZE);
+
+  // Use DataView to set the length of the payload.
+  const view = new DataView(frame.buffer, frame.byteOffset, frame.byteLength);
+  view.setUint32(1, data.length, false /* big endian */);
+
+  // Set the type in the MSB of the first byte.
+  frame[0] = type;
+
+  // Copy the bytes of the data into the frame.
+  frame.set(data, HEADER_SIZE);
+
+  return frame;
 }
 
-const SERVER = "http://alpha.web3portal.com:8081";
+export function serializeBrowserHeaders(headers: Metadata): string {
+  let result = "";
+  headers.forEach((key, values) => {
+    values.forEach((value) => {
+      result += `${key}: ${value}\r\n`;
+    });
+  });
+  return result;
+}
+
+export function encodeASCII(input: string): Uint8Array {
+  const encoded = new Uint8Array(input.length);
+  for (let i = 0; i !== input.length; ++i) {
+    const charCode = input.charCodeAt(i);
+    if (!isValidHeaderAscii(charCode)) {
+      throw new Error("Metadata contains invalid ASCII");
+    }
+    encoded[i] = charCode;
+  }
+  return encoded;
+}
+
+const isAllowedControlChars = (char: number) =>
+  char === 0x9 || char === 0xa || char === 0xd;
+
+function isValidHeaderAscii(val: number): boolean {
+  return isAllowedControlChars(val) || (val >= 0x20 && val <= 0x7e);
+}
+
+function invoke(methodDescriptor: any, props: any) {
+  const grpcClient = client(methodDescriptor, {
+    host: props.host,
+    transport: props.transport,
+    debug: props.debug,
+  });
+
+  if (props.onHeaders) {
+    grpcClient.onHeaders(props.onHeaders);
+  }
+  if (props.onMessage) {
+    grpcClient.onMessage(props.onMessage);
+  }
+  if (props.onEnd) {
+    grpcClient.onEnd(props.onEnd);
+  }
+
+  grpcClient.start(props.metadata);
+  grpcClient.send(props.request);
+  grpcClient.finishSend();
+
+  return {
+    client: grpcClient,
+    close: () => {
+      grpcClient.close();
+    },
+  };
+}
+
+const SERVER = "http://relay1.lumeweb.com:8082";
 const PROJECT_ID = "f195d68175eb091ec1f71d00f8952b85";
 const plugin = {
   name: "lavanet",
@@ -67,21 +138,54 @@ const plugin = {
         if (!req.data) {
           throw new Error("invalid data");
         }
-        const requestPromise = new Promise<Uint8Array>((resolve, reject) => {
+        const requestPromise = new Promise<Uint8Array[]>((resolve, reject) => {
           const request =
             BadgeGenerator.GenerateBadge.requestType.deserializeBinary(
               req.data,
             );
           request.setProjectId(PROJECT_ID);
-          grpc.invoke(BadgeGenerator.GenerateBadge, {
+
+          let responses: Uint8Array[] = [];
+          let grpcHeaders: Metadata[] = [];
+
+          const grpcReq = invoke(BadgeGenerator.GenerateBadge, {
             request,
             host: SERVER,
             transport: transport,
+            onHeaders(headers: Metadata) {
+              responses.push(
+                frameRequest(
+                  encodeASCII(serializeBrowserHeaders(headers)),
+                  0x80,
+                ),
+              );
+
+              if (!headers.has("grpc-status")) {
+                // @ts-ignore
+                grpcReq.client.responseHeaders = undefined;
+                grpcHeaders.push(headers);
+              } else {
+                // @ts-ignore
+                grpcReq.client.responseTrailers = headers;
+
+                const newHeaders = new Metadata();
+
+                grpcHeaders.forEach((metadata) => {
+                  metadata.forEach((key, values: string[]) => {
+                    newHeaders.append(key, values);
+                  });
+                });
+
+                // @ts-ignore
+                grpcReq.client.responseHeaders = newHeaders;
+              }
+            },
             onMessage(message: GenerateBadgeResponse) {
-              resolve(frameRequest(message));
+              responses.push(frameRequest(message.serializeBinary()));
             },
             onEnd(code, msg) {
               if (code === grpc.Code.OK || msg === undefined) {
+                resolve(responses);
                 return;
               }
               reject(
@@ -93,7 +197,12 @@ const plugin = {
             },
           });
         });
-        return relayWithTimeout<Uint8Array>(5000, requestPromise);
+        const ret = await relayWithTimeout<Uint8Array[]>(5000, requestPromise);
+        if (ret instanceof Error) {
+          return ret;
+        }
+
+        return ret.map((item) => Array.from(item));
       },
     });
   },