DEV: Fixes for modifyClass (#28227)

This resolves issues when a mix of callback-based modifications and Ember-reopen-based modifications are used on the same target. In summary:

- Fixes `pluginId` exception logic for callback-based modifications

- Moves `pluginId` storage to a WeakMap so it doesn't pollute the target's descriptors

- When applying a legacy modifyClass, we will temporarily rollback any modern callback-based modifications. This means all of Ember's reopen calls apply to un-prepended classes, and then we add our modern prepends on top.

- Calls `.proto()` on CoreObject descendants before prepending, to ensure that pending Ember mixins have been applied

app/assets/javascripts/discourse/app/lib/class-prepend.js

@@ -1,9 +1,19 @@
-import { DEBUG } from "@glimmer/env";
-import { isTesting } from "discourse-common/config/environment";
+import CoreObject from "@ember/object/core";
 
 const RESERVED_CLASS_PROPS = ["prototype", "name", "length"];
 const RESERVED_PROTOTYPE_PROPS = ["constructor"];
 
+function hasAncestor(klass, ancestor) {
+  let current = klass;
+  while (current) {
+    if (current === ancestor) {
+      return true;
+    }
+    current = Object.getPrototypeOf(current);
+  }
+  return false;
+}
+
 /**
  * This function provides a way to add/modify instance and static properties on an existing JS class, including
  * the ability to use `super` to call the original implementation.
@@ -13,9 +23,14 @@ const RESERVED_PROTOTYPE_PROPS = ["constructor"];
  *
  */
 export default function classPrepend(klass, callback) {
+  if (hasAncestor(klass, CoreObject)) {
+    // Ensure any prior reopen() calls have been applied
+    klass.proto();
+  }
+
   const originalKlassDescs = Object.getOwnPropertyDescriptors(klass);
   const originalProtoDescs = Object.getOwnPropertyDescriptors(klass.prototype);
-  logDescriptorInfoForRollback(klass, originalKlassDescs, originalProtoDescs);
+  logInfo(klass, originalKlassDescs, originalProtoDescs, callback);
 
   for (const key of RESERVED_CLASS_PROPS) {
     delete originalKlassDescs[key];
@@ -64,36 +79,73 @@ export default function classPrepend(klass, callback) {
   }
 }
 
-let originalDescriptorInfo;
+const prependInfo = new Map();
 
-if (DEBUG && isTesting()) {
-  originalDescriptorInfo = new Map();
-}
-
-function logDescriptorInfoForRollback(klass, klassDescs, protoDescs) {
-  if (DEBUG && isTesting() && !originalDescriptorInfo.has(klass)) {
-    originalDescriptorInfo.set(klass, {
+/**
+ * Log the previous state of a class so that it can be rolled back later.
+ */
+function logInfo(klass, klassDescs, protoDescs, modifyCallback) {
+  const info = prependInfo.get(klass) || {
     klassDescs,
     protoDescs,
-    });
+    modifyCallbacks: [],
+  };
+  info.modifyCallbacks.push(modifyCallback);
+  prependInfo.set(klass, info);
+}
+
+/**
+ * Rollback a specific class to its state before any prepends were applied.
+ */
+function rollbackPrepends(klass) {
+  const { klassDescs, protoDescs } = prependInfo.get(klass);
+
+  for (const [key, descriptor] of Object.entries(klassDescs)) {
+    Object.defineProperty(klass, key, descriptor);
+  }
+
+  for (const key of Object.getOwnPropertyNames(klass)) {
+    if (!RESERVED_CLASS_PROPS.includes(key) && !klassDescs[key]) {
+      delete klass[key];
+    }
+  }
+
+  for (const [key, descriptor] of Object.entries(protoDescs)) {
+    Object.defineProperty(klass.prototype, key, descriptor);
+  }
+
+  for (const key of Object.getOwnPropertyNames(klass.prototype)) {
+    if (!RESERVED_PROTOTYPE_PROPS.includes(key) && !protoDescs[key]) {
+      delete klass.prototype[key];
+    }
+  }
+
+  prependInfo.delete(klass);
+}
+
+/**
+ * Rollback all prepends on a class, run a callback, then re-apply the prepends.
+ */
+export function withPrependsRolledBack(klass, callback) {
+  const info = prependInfo.get(klass);
+  if (!info) {
+    callback();
+    return;
+  }
+
+  rollbackPrepends(klass);
+  try {
+    callback();
+  } finally {
+    info.modifyCallbacks.forEach((cb) => classPrepend(klass, cb));
   }
 }
 
 /**
  * Rollback all descriptors to their original values. This should only be used in tests
  */
-export function rollbackAllModifications() {
-  if (DEBUG && isTesting()) {
-    for (const [klass, { klassDescs, protoDescs }] of originalDescriptorInfo) {
-      for (const [key, descriptor] of Object.entries(klassDescs)) {
-        Object.defineProperty(klass, key, descriptor);
-      }
-
-      for (const [key, descriptor] of Object.entries(protoDescs)) {
-        Object.defineProperty(klass.prototype, key, descriptor);
-      }
-    }
-
-    originalDescriptorInfo.clear();
+export function rollbackAllPrepends() {
+  for (const klass of prependInfo.keys()) {
+    rollbackPrepends(klass);
   }
 }

app/assets/javascripts/discourse/app/lib/plugin-api.gjs

@@ -61,7 +61,9 @@ import {
   PLUGIN_NAV_MODE_TOP,
   registerAdminPluginConfigNav,
 } from "discourse/lib/admin-plugin-config-nav";
-import classPrepend from "discourse/lib/class-prepend";
+import classPrepend, {
+  withPrependsRolledBack,
+} from "discourse/lib/class-prepend";
 import { addPopupMenuOption } from "discourse/lib/composer/custom-popup-menu-options";
 import { registerDesktopNotificationHandler } from "discourse/lib/desktop-notifications";
 import { downloadCalendar } from "discourse/lib/download-calendar";
@@ -175,9 +177,15 @@ const DEPRECATED_HEADER_WIDGETS = [
   "user-dropdown",
 ];
 
+const appliedModificationIds = new WeakMap();
+
 // This helper prevents us from applying the same `modifyClass` over and over in test mode.
 function canModify(klass, type, resolverName, changes) {
-  if (typeof changes !== "function" && !changes.pluginId) {
+  if (typeof changes === "function") {
+    return true;
+  }
+
+  if (!changes.pluginId) {
     // eslint-disable-next-line no-console
     console.warn(
       consolePrefix(),
@@ -187,10 +195,13 @@ function canModify(klass, type, resolverName, changes) {
   }
 
   let key = "_" + type + "/" + changes.pluginId + "/" + resolverName;
-  if (klass.class[key]) {
+
+  if (appliedModificationIds.get(klass.class)?.includes(key)) {
     return false;
   } else {
-    klass.class[key] = 1;
+    const modificationIds = appliedModificationIds.get(klass.class) || [];
+    modificationIds.push(key);
+    appliedModificationIds.set(klass.class, modificationIds);
     return true;
   }
 }
@@ -299,7 +310,9 @@ class PluginApi {
       if (typeof changes === "function") {
         classPrepend(klass.class, changes);
       } else if (klass.class.reopen) {
+        withPrependsRolledBack(klass.class, () => {
           klass.class.reopen(changes);
+        });
       } else {
         Object.defineProperties(
           klass.class.prototype || klass.class,
@@ -330,7 +343,9 @@ class PluginApi {
 
     if (canModify(klass, "static", resolverName, changes)) {
       delete changes.pluginId;
+      withPrependsRolledBack(klass.class, () => {
         klass.class.reopenClass(changes);
+      });
     }
 
     return klass;

app/assets/javascripts/discourse/tests/helpers/qunit-helpers.js

@@ -32,7 +32,7 @@ import { clearHTMLCache } from "discourse/helpers/custom-html";
 import { resetUsernameDecorators } from "discourse/helpers/decorate-username-selector";
 import { resetBeforeAuthCompleteCallbacks } from "discourse/instance-initializers/auth-complete";
 import { resetAdminPluginConfigNav } from "discourse/lib/admin-plugin-config-nav";
-import { rollbackAllModifications } from "discourse/lib/class-prepend";
+import { rollbackAllPrepends } from "discourse/lib/class-prepend";
 import { clearPopupMenuOptions } from "discourse/lib/composer/custom-popup-menu-options";
 import { clearDesktopNotificationHandlers } from "discourse/lib/desktop-notifications";
 import { cleanUpHashtagTypeClasses } from "discourse/lib/hashtag-type-registry";
@@ -250,7 +250,7 @@ export function testCleanup(container, app) {
   clearAdditionalAdminSidebarSectionLinks();
   resetAdminPluginConfigNav();
   resetTransformers();
-  rollbackAllModifications();
+  rollbackAllPrepends();
 }
 
 function cleanupCssGeneratorTags() {

app/assets/javascripts/discourse/tests/unit/lib/class-prepend-test.js

@@ -1,9 +1,7 @@
 import { tracked } from "@glimmer/tracking";
 import { action } from "@ember/object";
 import { module, test } from "qunit";
-import classPrepend, {
-  rollbackAllModifications,
-} from "discourse/lib/class-prepend";
+import classPrepend, { rollbackAllPrepends } from "discourse/lib/class-prepend";
 
 module("Unit | class-prepend", function () {
   test("can override function, with super support", function (assert) {
@@ -273,7 +271,7 @@ module("Unit | class-prepend", function () {
 
     assert.strictEqual(new Topic().someFunction(), 2, "change is applied");
 
-    rollbackAllModifications();
+    rollbackAllPrepends();
 
     assert.strictEqual(new Topic().someFunction(), 1, "change is rolled back");
   });

app/assets/javascripts/discourse/tests/unit/lib/plugin-api-test.js

@@ -2,6 +2,7 @@ import EmberObject from "@ember/object";
 import { getOwner } from "@ember/owner";
 import { setupTest } from "ember-qunit";
 import { module, test } from "qunit";
+import { rollbackAllPrepends } from "discourse/lib/class-prepend";
 import { withPluginApi } from "discourse/lib/plugin-api";
 import discourseComputed from "discourse-common/utils/decorators";
 
@@ -153,10 +154,23 @@ module("Unit | Utility | plugin-api", function (hooks) {
           }
       );
 
+      api.modifyClass(
+        "test-thingy:main",
+        (Superclass) =>
+          class extends Superclass {
+            someFunction() {
+              return `${super.someFunction()} twice`;
+            }
+          }
+      );
+
       const thingyKlass =
         getOwner(this).resolveRegistration("test-thingy:main");
       const thingy = new thingyKlass();
-      assert.strictEqual(thingy.someFunction(), "original function modified");
+      assert.strictEqual(
+        thingy.someFunction(),
+        "original function modified twice"
+      );
       assert.strictEqual(thingy.someGetter, "original getter modified");
       assert.strictEqual(
         TestThingy.someStaticMethod(),
@@ -164,4 +178,61 @@ module("Unit | Utility | plugin-api", function (hooks) {
       );
     });
   });
+
+  test("modifyClass works with a combination of callback and legacy syntax", function (assert) {
+    class TestThingy extends EmberObject {
+      someMethod() {
+        return "original";
+      }
+    }
+
+    getOwner(this).register("test-thingy:main", TestThingy);
+
+    const fakeInit = () => {
+      withPluginApi("1.1.0", (api) => {
+        api.modifyClass("test-thingy:main", {
+          someMethod() {
+            return `${this._super()} reopened`;
+          },
+          pluginId: "one",
+        });
+
+        api.modifyClass(
+          "test-thingy:main",
+          (Superclass) =>
+            class extends Superclass {
+              someMethod() {
+                return `${super.someMethod()}, prepended`;
+              }
+            }
+        );
+
+        api.modifyClass("test-thingy:main", {
+          someMethod() {
+            return `${this._super()}, reopened2`;
+          },
+          pluginId: "two",
+        });
+      });
+    };
+
+    fakeInit();
+
+    assert.strictEqual(
+      new TestThingy().someMethod(),
+      "original reopened, reopened2, prepended",
+      "it works after first application"
+    );
+
+    for (let i = 0; i < 3; i++) {
+      rollbackAllPrepends();
+      fakeInit();
+    }
+
+    assert.strictEqual(
+      new TestThingy().someMethod(),
+      "original reopened, reopened2, prepended",
+      "it works when rolled back and re-applied multiple times"
+    );
+  });
 });