lib: make IterableWeakMap safe to iterate

aduh95 · targos · commit cea8b4265c5b · 2021-06-11T09:30:14.000+02:00
PR-URL: #38523 Reviewed-By: Rich Trott <rtrott@gmail.com> Reviewed-By: James M Snell <jasnell@gmail.com>
diff --git a/lib/internal/util/iterable_weak_map.js b/lib/internal/util/iterable_weak_map.js
@@ -65,13 +65,22 @@ class IterableWeakMap {
     return true;
   }
 
-  *[SymbolIterator]() {
-    for (const ref of this.#refSet) {
-      const key = ref.deref();
-      if (!key) continue;
+  [SymbolIterator]() {
+    const iterator = this.#refSet[SymbolIterator]();
+
+    const next = () => {
+      const result = iterator.next();
+      if (result.done) return result;
+      const key = result.value.deref();
+      if (key == null) return next();
       const { value } = this.#weakMap.get(key);
-      yield value;
-    }
+      return { done: false, value };
+    };
+
+    return {
+      [SymbolIterator]() { return this; },
+      next,
+    };
   }
 }
 
diff --git a/test/parallel/test-internal-iterable-weak-map.js b/test/parallel/test-internal-iterable-weak-map.js
@@ -1,10 +1,16 @@
 // Flags: --expose-gc --expose-internals
 'use strict';
 
-require('../common');
+const common = require('../common');
 const { deepStrictEqual, strictEqual } = require('assert');
 const { IterableWeakMap } = require('internal/util/iterable_weak_map');
 
+// Ensures iterating over the map does not rely on methods which can be
+// mutated by users.
+Reflect.getPrototypeOf(function*() {}).prototype.next = common.mustNotCall();
+Reflect.getPrototypeOf(new Set()[Symbol.iterator]()).next =
+  common.mustNotCall();
+
 // It drops entry if a reference is no longer held.
 {
   const wm = new IterableWeakMap();
