From 52bc66e7dd934c22913a562757048ec48bde0e58 Mon Sep 17 00:00:00 2001
From: Jorrin <jorrinkievit@hotmail.com>
Date: Tue, 9 Jan 2024 23:00:54 +0100
Subject: [PATCH] improve typings

---
 src/backend/extension/messaging.ts | 51 ++++++++++++++---------
 src/backend/extension/plasmo.ts    | 67 ++++++++++++++++++++----------
 src/backend/extension/streams.ts   |  9 ++--
 src/backend/providers/fetchers.ts  |  9 +++-
 4 files changed, 87 insertions(+), 49 deletions(-)

diff --git a/src/backend/extension/messaging.ts b/src/backend/extension/messaging.ts
index 3c5480c1..28f80ad2 100644
--- a/src/backend/extension/messaging.ts
+++ b/src/backend/extension/messaging.ts
@@ -7,18 +7,17 @@ import { isAllowedExtensionVersion } from "@/backend/extension/compatibility";
 
 let activeExtension = false;
 
-export interface ExtensionHello {
-  version: string;
-}
-
-function sendMessage<T, Payload>(
-  message: keyof MessagesMetadata,
-  payload: any,
+function sendMessage<MessageKey extends keyof MessagesMetadata>(
+  message: MessageKey,
+  payload: MessagesMetadata[MessageKey]["req"],
   timeout: number = -1,
 ) {
-  return new Promise<T | null>((resolve) => {
+  return new Promise<MessagesMetadata[MessageKey]["res"] | null>((resolve) => {
     if (timeout >= 0) setTimeout(() => resolve(null), timeout);
-    sendToBackgroundViaRelay<Payload, T>({
+    sendToBackgroundViaRelay<
+      MessagesMetadata[MessageKey]["req"],
+      MessagesMetadata[MessageKey]["res"]
+    >({
       name: message,
       body: payload,
     })
@@ -34,21 +33,33 @@ function sendMessage<T, Payload>(
 }
 
 export async function sendExtensionRequest(
-  url: string,
-  ops: any,
-): Promise<ExtensionHello | null> {
-  return sendMessage("make-request", { url, ...ops });
+  ops: Omit<MessagesMetadata["makeRequest"]["req"], "requestDomain">,
+): Promise<MessagesMetadata["makeRequest"]["res"] | null> {
+  return sendMessage("makeRequest", {
+    requestDomain: window.location.origin,
+    ...ops,
+  });
 }
 
 export async function setDomainRule(
-  domains: string[],
-  headers: Record<string, string>,
-): Promise<ExtensionHello | null> {
-  return sendMessage("prepare-stream", { domains, headers });
+  ops: Omit<MessagesMetadata["prepareStream"]["req"], "requestDomain">,
+): Promise<MessagesMetadata["prepareStream"]["res"] | null> {
+  return sendMessage("prepareStream", {
+    requestDomain: window.location.origin,
+    ...ops,
+  });
 }
 
-export async function extensionInfo(): Promise<ExtensionHello | null> {
-  return sendMessage("hello", null, 300);
+export async function extensionInfo(): Promise<
+  MessagesMetadata["hello"]["res"] | null
+> {
+  return sendMessage(
+    "hello",
+    {
+      requestDomain: window.location.origin,
+    },
+    300,
+  );
 }
 
 export function isExtensionActiveCached(): boolean {
@@ -57,7 +68,7 @@ export function isExtensionActiveCached(): boolean {
 
 export async function isExtensionActive(): Promise<boolean> {
   const info = await extensionInfo();
-  if (!info) return false;
+  if (!info?.success) return false;
   const allowedVersion = isAllowedExtensionVersion(info.version);
   if (!allowedVersion) return false;
   return true;
diff --git a/src/backend/extension/plasmo.ts b/src/backend/extension/plasmo.ts
index 8c0093f0..6c37ad14 100644
--- a/src/backend/extension/plasmo.ts
+++ b/src/backend/extension/plasmo.ts
@@ -1,6 +1,37 @@
-export interface PlasmoRequestBody {
+export interface ExtensionBaseRequest {
+  requestDomain: string;
+}
+
+export type ExtensionBaseResponse<T = object> =
+  | ({
+      success: true;
+    } & T)
+  | {
+      success: false;
+      error: string;
+    };
+
+export type ExtensionHelloResponse = ExtensionBaseResponse<{
+  version: string;
+}>;
+
+export interface ExtensionMakeRequest extends ExtensionBaseRequest {
+  url: string;
+  method: string;
+  headers?: Record<string, string>;
+  body?: string | FormData | URLSearchParams | Record<string, any>;
+}
+
+export type ExtensionMakeRequestResponse = ExtensionBaseResponse<{
+  status: number;
+  requestHeaders: Record<string, string>;
+  responseHeaders: Record<string, string>;
+  data: string | Record<string, unknown>;
+}>;
+
+export interface ExtensionPrepareStreamRequest extends ExtensionBaseRequest {
   ruleId: number;
-  domain: string;
+  targetDomains: string[];
   requestHeaders?: Record<string, string>;
   responseHeaders?: Record<string, string>;
 }
@@ -9,28 +40,18 @@ export interface ExtensionHelloReply {
   version: string;
 }
 
-export type ExtensionRequestReply =
-  | {
-      success: true;
-      ruleId: number;
-    }
-  | {
-      success: false;
-      error: string;
-    };
-
-interface MmMetadata {
-  "prepare-stream": {
-    req: PlasmoRequestBody;
-    res: ExtensionRequestReply;
-  };
-  "make-request": {
-    req: PlasmoRequestBody;
-    res: ExtensionRequestReply;
-  };
+export interface MmMetadata {
   hello: {
-    req: null;
-    res: ExtensionHelloReply;
+    req: ExtensionBaseRequest;
+    res: ExtensionHelloResponse;
+  };
+  makeRequest: {
+    req: ExtensionMakeRequest;
+    res: ExtensionMakeRequestResponse;
+  };
+  prepareStream: {
+    req: ExtensionPrepareStreamRequest;
+    res: ExtensionBaseResponse;
   };
 }
 
diff --git a/src/backend/extension/streams.ts b/src/backend/extension/streams.ts
index 8d202992..2afb900d 100644
--- a/src/backend/extension/streams.ts
+++ b/src/backend/extension/streams.ts
@@ -33,8 +33,9 @@ function buildHeadersFromStream(stream: Stream): Record<string, string> {
 }
 
 export async function prepareStream(stream: Stream) {
-  await setDomainRule(
-    extractDomainsFromStream(stream),
-    buildHeadersFromStream(stream),
-  );
+  await setDomainRule({
+    ruleId: 1,
+    targetDomains: extractDomainsFromStream(stream),
+    requestHeaders: buildHeadersFromStream(stream),
+  });
 }
diff --git a/src/backend/providers/fetchers.ts b/src/backend/providers/fetchers.ts
index 596e8376..e133199c 100644
--- a/src/backend/providers/fetchers.ts
+++ b/src/backend/providers/fetchers.ts
@@ -54,8 +54,13 @@ export function makeLoadBalancedSimpleProxyFetcher() {
 }
 
 export function makeExtensionFetcher() {
-  const fetcher: Fetcher = async (a, b) => {
-    return sendExtensionRequest(a, b) as any;
+  const fetcher: Fetcher = async (url, ops) => {
+    return sendExtensionRequest({
+      url,
+      method: ops?.method ?? "GET",
+      headers: ops?.headers,
+      body: ops?.body,
+    }) as any;
   };
   return fetcher;
 }