From cea8b4265c5bfd44f2400fcc546036bd58d058c0 Mon Sep 17 00:00:00 2001
From: Antoine du Hamel <duhamelantoine1995@gmail.com>
Date: Mon, 3 May 2021 18:32:14 +0200
Subject: [PATCH] lib: make `IterableWeakMap` safe to iterate

PR-URL: https://github.com/nodejs/node/pull/38523
Reviewed-By: Rich Trott <rtrott@gmail.com>
Reviewed-By: James M Snell <jasnell@gmail.com>
---
 lib/internal/util/iterable_weak_map.js        | 21 +++++++++++++------
 .../test-internal-iterable-weak-map.js        |  8 ++++++-
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/lib/internal/util/iterable_weak_map.js b/lib/internal/util/iterable_weak_map.js
index c9715a7e313b20..f82d41b1f962ba 100644
--- 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
index e0282c9081ee33..f2befe13da87f3 100644
--- 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();