From 326af1930faab421bdd5290fdb219b31246bbd72 Mon Sep 17 00:00:00 2001
From: GarboMuffin <muffin@mailbox.org>
Date: Sat, 22 Oct 2022 16:03:05 -0500
Subject: [PATCH] Clean up unsandboxed extension loader

Remove a bit of global state hackery and prevent concurrent loads
---
 src/extension-support/extension-manager.js    | 36 +++++++++++-----
 .../tw-unsandboxed-extension-runner.js        | 41 ++++++++++---------
 2 files changed, 46 insertions(+), 31 deletions(-)

diff --git a/src/extension-support/extension-manager.js b/src/extension-support/extension-manager.js
index 25fa59bc..8d17fed5 100644
--- a/src/extension-support/extension-manager.js
+++ b/src/extension-support/extension-manager.js
@@ -92,8 +92,14 @@ class ExtensionManager {
         this.pendingWorkers = [];
 
         /**
-         * Set of loaded extension URLs/IDs (equivalent for built-in extensions).
-         * @type {Set.<string>}
+         * Map of worker ID to the URL where it was loaded from.
+         * @type {Array<string>}
+         */
+        this.workerURLs = [];
+
+        /**
+         * Map of loaded extension URLs/IDs to service names.
+         * @type {Map.<string, string>}
          * @private
          */
         this._loadedExtensions = new Map();
@@ -185,13 +191,21 @@ class ExtensionManager {
         this.loadingAsyncExtensions++;
 
         if (this.workerMode === 'unsandboxed') {
-            return require('./tw-unsandboxed-extension-runner')
-                .load(
-                    extensionURL,
-                    this._registerInternalExtension.bind(this)
-                )
-                .then(() => {
-                    this._finishedLoadingExtension();
+            const {load} = require('./tw-unsandboxed-extension-runner');
+            return load(extensionURL)
+                .then(extensionObjects => {
+                    const fakeWorkerId = this.nextExtensionWorker++;
+                    this.workerURLs[fakeWorkerId] = extensionURL;
+
+                    for (const extensionObject of extensionObjects) {
+                        const extensionInfo = extensionObject.getInfo();
+                        const serviceName = `unsandboxed.${fakeWorkerId}.${extensionInfo.id}`;
+                        dispatch.setServiceSync(serviceName, extensionObject);
+                        dispatch.callSync('extensions', 'registerExtensionServiceSync', serviceName);
+                        this._loadedExtensions.set(extensionInfo.id, serviceName);
+                    }
+
+                    this._finishedLoadingExtensionScript();
                 });
         }
 
@@ -274,11 +288,11 @@ class ExtensionManager {
         dispatch.call(serviceName, 'getInfo').then(info => {
             this._loadedExtensions.set(info.id, serviceName);
             this._registerExtensionInfo(serviceName, info);
-            this._finishedLoadingExtension();
+            this._finishedLoadingExtensionScript();
         });
     }
 
-    _finishedLoadingExtension () {
+    _finishedLoadingExtensionScript () {
         this.loadingAsyncExtensions--;
         if (this.loadingAsyncExtensions === 0) {
             this.asyncExtensionsLoadedCallbacks.forEach(i => i());
diff --git a/src/extension-support/tw-unsandboxed-extension-runner.js b/src/extension-support/tw-unsandboxed-extension-runner.js
index 64bd3e50..e82efe32 100644
--- a/src/extension-support/tw-unsandboxed-extension-runner.js
+++ b/src/extension-support/tw-unsandboxed-extension-runner.js
@@ -1,37 +1,38 @@
 const ArgumentType = require('../extension-support/argument-type');
 const BlockType = require('../extension-support/block-type');
 const TargetType = require('../extension-support/target-type');
+const AsyncLimiter = require('../util/async-limiter');
 
-// Custom extensions by nature rely on the single global Scratch object to register themselves
-// To make this work, we'll have to track the most recent callback and just hope that we don't
-// have two VMs on the same page trying to load unsandboxed extensions.
-let mostRecentRegisterExtensionCallback = null;
+const loadUnsandboxedExtension = extensionURL => new Promise((resolve, reject) => {
+    global.Scratch = global.Scratch || {};
+    global.Scratch.ArgumentType = ArgumentType;
+    global.Scratch.BlockType = BlockType;
+    global.Scratch.TargetType = TargetType;
+    global.Scratch.ArgumentType = ArgumentType;
 
-const register = extensionObject => {
-    mostRecentRegisterExtensionCallback(extensionObject);
-};
+    // When we load the script, wait for it to eventually call Scratch.extensions.register
+    // before resolving. Note that extension scripts may register multiple extensions.
 
-const Scratch = {
-    ArgumentType,
-    BlockType,
-    TargetType,
-    extensions: {
+    const extensionObjects = [];
+    const register = extensionObject => {
+        extensionObjects.push(extensionObject);
+        resolve(extensionObjects);
+    };
+    global.Scratch.extensions = {
         register
-    }
-};
-
-global.Scratch = Scratch;
-
-const load = (extensionURL, registerExtensionCallback) => new Promise((resolve, reject) => {
-    mostRecentRegisterExtensionCallback = registerExtensionCallback;
+    };
 
     const script = document.createElement('script');
-    script.onload = () => resolve();
     script.onerror = () => reject(new Error(`Error in unsandboxed script ${extensionURL}`));
     script.src = extensionURL;
     document.body.appendChild(script);
 });
 
+// Because loading unsandboxed extensions requires messing with global state (global.Scratch),
+// only let one extension load at a time.
+const limiter = new AsyncLimiter(loadUnsandboxedExtension, 1);
+const load = extensionURL => limiter.do(extensionURL);
+
 module.exports = {
     load
 };