From d882f64daa9ddb55a5a52afc4421b2da076fe40d Mon Sep 17 00:00:00 2001
From: Frederik Bolding <frederik.bolding@gmail.com>
Date: Tue, 3 Dec 2024 12:35:02 +0100
Subject: [PATCH 1/2] feat: Allow passing btoa for mobile execution service

---
 .../webview/WebViewExecutionService.ts        |  3 +++
 .../services/webview/WebViewMessageStream.ts  | 20 +++++++++++++++----
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/packages/snaps-controllers/src/services/webview/WebViewExecutionService.ts b/packages/snaps-controllers/src/services/webview/WebViewExecutionService.ts
index 8f5a9166c6..b546afcfe0 100644
--- a/packages/snaps-controllers/src/services/webview/WebViewExecutionService.ts
+++ b/packages/snaps-controllers/src/services/webview/WebViewExecutionService.ts
@@ -5,6 +5,7 @@ import { WebViewMessageStream } from './WebViewMessageStream';
 
 export type WebViewExecutionServiceArgs = ExecutionServiceArgs & {
   getWebView: () => Promise<WebViewInterface>;
+  btoa?: (data: string) => string;
 };
 
 export class WebViewExecutionService extends ProxyExecutionService {
@@ -14,6 +15,7 @@ export class WebViewExecutionService extends ProxyExecutionService {
     messenger,
     setupSnapProvider,
     getWebView,
+    btoa,
   }: WebViewExecutionServiceArgs) {
     super({
       messenger,
@@ -22,6 +24,7 @@ export class WebViewExecutionService extends ProxyExecutionService {
         name: 'parent',
         target: 'child',
         getWebView,
+        btoa,
       }),
     });
     this.#getWebView = getWebView;
diff --git a/packages/snaps-controllers/src/services/webview/WebViewMessageStream.ts b/packages/snaps-controllers/src/services/webview/WebViewMessageStream.ts
index b8720353a5..662dbfeba3 100644
--- a/packages/snaps-controllers/src/services/webview/WebViewMessageStream.ts
+++ b/packages/snaps-controllers/src/services/webview/WebViewMessageStream.ts
@@ -14,6 +14,7 @@ export type WebViewStreamArgs = {
   name: string;
   target: string;
   getWebView: () => Promise<WebViewInterface>;
+  btoa?: (data: string) => string;
 };
 
 /**
@@ -26,6 +27,8 @@ export class WebViewMessageStream extends BasePostMessageStream {
 
   #webView: WebViewInterface | undefined;
 
+  #btoa?: (data: string) => string;
+
   /**
    * Creates a stream for communicating with other streams inside a WebView.
    *
@@ -34,12 +37,14 @@ export class WebViewMessageStream extends BasePostMessageStream {
    * multiple streams sharing the same window object.
    * @param args.target - The name of the stream to exchange messages with.
    * @param args.getWebView - A asynchronous getter for the webview.
+   * @param args.btoa - An optional function that encodes a string to base64.
    */
-  constructor({ name, target, getWebView }: WebViewStreamArgs) {
+  constructor({ name, target, getWebView, btoa }: WebViewStreamArgs) {
     super();
 
     this.#name = name;
     this.#target = target;
+    this.#btoa = btoa;
 
     this._onMessage = this._onMessage.bind(this);
 
@@ -58,6 +63,15 @@ export class WebViewMessageStream extends BasePostMessageStream {
       });
   }
 
+  #encodeMessage(json: string): string {
+    if (this.#btoa) {
+      return this.#btoa(json);
+    }
+
+    const bytes = stringToBytes(json);
+    return bytesToBase64(bytes);
+  }
+
   protected _postMessage(data: unknown): void {
     assert(this.#webView);
     const json = JSON.stringify({
@@ -67,9 +81,7 @@ export class WebViewMessageStream extends BasePostMessageStream {
 
     // To prevent XSS, we base64 encode the message before injecting it.
     // This adds significant performance overhead.
-    // TODO: Should we use mobile native base64 here?
-    const bytes = stringToBytes(json);
-    const base64 = bytesToBase64(bytes);
+    const base64 = this.#encodeMessage(json);
     this.#webView.injectJavaScript(`window.postMessage('${base64}')`);
   }
 

From ae074dbc091ce51f2b3e2351fe8bac140f1325eb Mon Sep 17 00:00:00 2001
From: Frederik Bolding <frederik.bolding@gmail.com>
Date: Tue, 3 Dec 2024 12:40:31 +0100
Subject: [PATCH 2/2] Add test

---
 packages/snaps-controllers/coverage.json             |  4 ++--
 packages/snaps-controllers/src/test-utils/webview.ts | 12 +++++++++++-
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/packages/snaps-controllers/coverage.json b/packages/snaps-controllers/coverage.json
index 28873f6536..0ef4e9b911 100644
--- a/packages/snaps-controllers/coverage.json
+++ b/packages/snaps-controllers/coverage.json
@@ -1,6 +1,6 @@
 {
-  "branches": 92.89,
+  "branches": 92.91,
   "functions": 96.71,
-  "lines": 98,
+  "lines": 98.01,
   "statements": 97.71
 }
diff --git a/packages/snaps-controllers/src/test-utils/webview.ts b/packages/snaps-controllers/src/test-utils/webview.ts
index 4838f57c95..2e9d1a9fe1 100644
--- a/packages/snaps-controllers/src/test-utils/webview.ts
+++ b/packages/snaps-controllers/src/test-utils/webview.ts
@@ -1,4 +1,9 @@
-import { base64ToBytes, bytesToString } from '@metamask/utils';
+import {
+  base64ToBytes,
+  bytesToBase64,
+  bytesToString,
+  stringToBytes,
+} from '@metamask/utils';
 
 import { WebViewMessageStream } from '../services/webview/WebViewMessageStream';
 
@@ -52,6 +57,11 @@ export function createWebViewObjects() {
     name: 'a',
     target: 'b',
     getWebView: mockGetWebViewA,
+    // For one of the streams, we test that a custom btoa function is used.
+    btoa: (data: string) => {
+      const bytes = stringToBytes(data);
+      return bytesToBase64(bytes);
+    },
   });
 
   const streamB = new WebViewMessageStream({