From f70a75a905b31c2c79123dd43acc70dbc8a2aa59 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Florian=20M=C3=BCllner?= <fmuellner@gnome.org>
Date: Fri, 23 Jun 2023 19:53:05 +0200
Subject: [PATCH] extensionUtils: Add InjectionManager

It is fairly common for extensions to monkey-patch existing
classes. Add a small helper class that makes this a tad bit
more convenient.

Part-of: <https://gitlab.gnome.org/GNOME/gnome-shell/-/merge_requests/2809>
---
 js/extensions/extension.js     | 68 +++++++++++++++++++++++++
 tests/meson.build              |  1 +
 tests/unit/injectionManager.js | 93 ++++++++++++++++++++++++++++++++++
 3 files changed, 162 insertions(+)
 create mode 100644 tests/unit/injectionManager.js

diff --git a/js/extensions/extension.js b/js/extensions/extension.js
index 2e06b3dfe..a0d3f049d 100644
--- a/js/extensions/extension.js
+++ b/js/extensions/extension.js
@@ -29,3 +29,71 @@ export class Extension extends ExtensionBase {
 export const {
     gettext, ngettext, pgettext,
 } = Extension.defineTranslationFunctions();
+
+export class InjectionManager {
+    #savedMethods = new Map();
+
+    /**
+     * @callback CreateOverrideFunc
+     * @param {Function?} originalMethod - the original method if it exists
+     * @returns {Function} - a function to be used as override
+     */
+
+    /**
+     * Modify, replace or inject a method
+     *
+     * @param {object} prototype - the object (or prototype) that is modified
+     * @param {string} methodName - the name of the overwritten method
+     * @param {CreateOverrideFunc} createOverrideFunc - function to call to create the override
+     */
+    overrideMethod(prototype, methodName, createOverrideFunc) {
+        const originalMethod = this._saveMethod(prototype, methodName);
+        prototype[methodName] = createOverrideFunc(originalMethod);
+    }
+
+    /**
+     * Restore the original method
+     *
+     * @param {object} prototype - the object (or prototype) that is modified
+     * @param {string} methodName - the name of the method to restore
+     */
+    restoreMethod(prototype, methodName) {
+        const savedProtoMethods = this.#savedMethods.get(prototype);
+        if (!savedProtoMethods)
+            return;
+
+        const originalMethod = savedProtoMethods.get(methodName);
+        if (originalMethod === undefined)
+            delete prototype[methodName];
+        else
+            prototype[methodName] = originalMethod;
+
+        savedProtoMethods.delete(methodName);
+        if (savedProtoMethods.size === 0)
+            this.#savedMethods.delete(prototype);
+    }
+
+    /**
+     * Restore all original methods and clear overrides
+     */
+    clear() {
+        for (const [proto, map] of this.#savedMethods) {
+            map.forEach(
+                (_, methodName) => this.restoreMethod(proto, methodName));
+        }
+        console.assert(this.#savedMethods.size === 0,
+            `${this.#savedMethods.size} overrides left after clear()`);
+    }
+
+    _saveMethod(prototype, methodName) {
+        let savedProtoMethods = this.#savedMethods.get(prototype);
+        if (!savedProtoMethods) {
+            savedProtoMethods = new Map();
+            this.#savedMethods.set(prototype, savedProtoMethods);
+        }
+
+        const originalMethod = prototype[methodName];
+        savedProtoMethods.set(methodName, originalMethod);
+        return originalMethod;
+    }
+}
diff --git a/tests/meson.build b/tests/meson.build
index df35080fe..a22a0d949 100644
--- a/tests/meson.build
+++ b/tests/meson.build
@@ -20,6 +20,7 @@ dbus_runner = configure_file(
 
 unit_tests = [
     'highlighter',
+    'injectionManager',
     'insertSorted',
     'jsParse',
     'markup',
diff --git a/tests/unit/injectionManager.js b/tests/unit/injectionManager.js
new file mode 100644
index 000000000..6de509441
--- /dev/null
+++ b/tests/unit/injectionManager.js
@@ -0,0 +1,93 @@
+const JsUnit = imports.jsUnit;
+
+import '../../js/ui/environment.js';
+
+import {InjectionManager} from '../../js/extensions/extension.js';
+
+class Object1 {
+    count = 0;
+
+    getNumber() {
+        return 42;
+    }
+
+    getCount() {
+        return ++this.count;
+    }
+}
+
+/**
+ * @param {object} object to modify
+ */
+function addInjections(object) {
+    // extend original method
+    injectionManager.overrideMethod(
+        object, 'getNumber', originalMethod => {
+            return function () {
+                // eslint-disable-next-line no-invalid-this
+                const num = originalMethod.call(this);
+                return 2 * num;
+            };
+        });
+
+    // override original method
+    injectionManager.overrideMethod(
+        object, 'getCount', () => {
+            return function () {
+                return 42;
+            };
+        });
+
+    // inject new method
+    injectionManager.overrideMethod(
+        object, 'getOtherNumber', () => {
+            return function () {
+                return 42;
+            };
+        });
+}
+
+
+const injectionManager = new InjectionManager();
+let obj;
+
+// Prototype injections
+addInjections(Object1.prototype);
+
+obj = new Object1();
+
+// new obj is modified
+JsUnit.assertEquals(obj.getNumber(), 84);
+JsUnit.assertEquals(obj.getCount(), 42);
+JsUnit.assertEquals(obj.count, 0);
+JsUnit.assertEquals(obj.getOtherNumber(), 42);
+
+injectionManager.clear();
+
+obj = new Object1();
+
+// new obj is unmodified
+JsUnit.assertEquals(obj.getNumber(), 42);
+JsUnit.assertEquals(obj.getCount(), obj.count);
+JsUnit.assert(obj.count > 0);
+JsUnit.assertRaises(() => obj.getOtherNumber());
+
+// instance injections
+addInjections(obj);
+
+// obj is now modified
+JsUnit.assertEquals(obj.getNumber(), 84);
+JsUnit.assertEquals(obj.getCount(), 42);
+JsUnit.assertEquals(obj.count, 1);
+JsUnit.assertEquals(obj.getOtherNumber(), 42);
+
+injectionManager.restoreMethod(obj, 'getNumber');
+JsUnit.assertEquals(obj.getNumber(), 42);
+
+injectionManager.clear();
+
+// obj is unmodified again
+JsUnit.assertEquals(obj.getNumber(), 42);
+JsUnit.assertEquals(obj.getCount(), obj.count);
+JsUnit.assert(obj.count > 0);
+JsUnit.assertRaises(() => obj.getOtherNumber());