From 204527d2691d65972cb9b8007a127d434d5da46b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Eduard=20Bardaj=C3=AD=20Puig?= <ebardaji@gmail.com>
Date: Fri, 26 Jul 2024 13:52:45 +0200
Subject: [PATCH] Remove RPC response source filtering

---
 src/content-scripts/content-script.ts |  6 ++--
 src/inpage/sats.inpage.ts             | 50 ++++++++++++++++++---------
 src/inpage/utils.ts                   |  5 ---
 3 files changed, 35 insertions(+), 26 deletions(-)

diff --git a/src/content-scripts/content-script.ts b/src/content-scripts/content-script.ts
index a82004548..bceeda6a2 100644
--- a/src/content-scripts/content-script.ts
+++ b/src/content-scripts/content-script.ts
@@ -64,10 +64,8 @@ function sendMessageToBackground(
 
 // Receives message from background script to execute in browser
 chrome.runtime.onMessage.addListener((message: LegacyMessageToContentScript) => {
-  if (message.source === MESSAGE_SOURCE) {
-    // Forward to web app (browser)
-    window.postMessage(message, window.location.origin);
-  }
+  // Forward to web app (browser)
+  window.postMessage(message, window.location.origin);
 });
 
 interface ForwardDomEventToBackgroundArgs {
diff --git a/src/inpage/sats.inpage.ts b/src/inpage/sats.inpage.ts
index d96c5ade8..f4b81a2c1 100644
--- a/src/inpage/sats.inpage.ts
+++ b/src/inpage/sats.inpage.ts
@@ -18,20 +18,22 @@ import {
   type SignMessageResponseMessage,
   type SignPsbtResponseMessage,
 } from '@common/types/message-types';
-import type {
-  BitcoinProvider,
-  CreateInscriptionResponse,
-  CreateRepeatInscriptionsResponse,
-  GetAddressResponse,
-  Params,
-  Requests,
-  RpcRequest,
-  RpcResponse,
-  SignMultipleTransactionsResponse,
-  SignTransactionResponse,
+import {
+  rpcResponseMessageSchema,
+  type BitcoinProvider,
+  type CreateInscriptionResponse,
+  type CreateRepeatInscriptionsResponse,
+  type GetAddressResponse,
+  type Params,
+  type Requests,
+  type RpcRequest,
+  type RpcResponse,
+  type SignMultipleTransactionsResponse,
+  type SignTransactionResponse,
 } from '@sats-connect/core';
 import { nanoid } from 'nanoid';
-import { isValidLegacyEvent, isValidRpcEvent } from './utils';
+import * as v from 'valibot';
+import { isValidLegacyEvent } from './utils';
 
 const SatsMethodsProvider: BitcoinProvider = {
   connect: async (btcAddressRequest): Promise<GetAddressResponse> => {
@@ -218,14 +220,28 @@ const SatsMethodsProvider: BitcoinProvider = {
     const rpcRequestEvent = new CustomEvent(DomEventName.rpcRequest, { detail: rpcRequest });
     document.dispatchEvent(rpcRequestEvent);
     return new Promise((resolve) => {
-      function handleRpcResponseEvent(eventMessage: MessageEvent<any>) {
-        if (!isValidRpcEvent(eventMessage)) return;
-        const response = eventMessage.data;
-        if (response.id !== id) {
+      function handleRpcResponseEvent(message: MessageEvent<unknown>) {
+        const parseResult = v.safeParse(rpcResponseMessageSchema, message.data);
+
+        if (!parseResult.success) {
+          // Ignore message if it's not an RPC message.
+          return;
+        }
+
+        const rpcResponseMessage = parseResult.output;
+
+        if (rpcResponseMessage.id !== id) {
+          // Ignore message if it's not a response to the current request.
           return;
         }
+
         window.removeEventListener('message', handleRpcResponseEvent);
-        return resolve(response);
+
+        // NOTE: Ideally the response would be runtime type-checked before the
+        // promise is resolved since the message crosses a type assertion
+        // boundary. For now, since all the responses are typed, it's relatively
+        // safe to assume that the message will conform to the expected type.
+        return resolve(rpcResponseMessage as RpcResponse<Method>);
       }
       window.addEventListener('message', handleRpcResponseEvent);
     });
diff --git a/src/inpage/utils.ts b/src/inpage/utils.ts
index 0001914cd..957c531f9 100644
--- a/src/inpage/utils.ts
+++ b/src/inpage/utils.ts
@@ -24,11 +24,6 @@ export const isValidLegacyEvent = (
   return correctSource && correctMethod && !!data.payload;
 };
 
-export const isValidRpcEvent = (event: MessageEvent) => {
-  const { data } = event;
-  return data.source === MESSAGE_SOURCE;
-};
-
 export const callAndReceive = async (
   methodName: CallableMethods | 'getURL',
   opts: any = {},