From 96624b71a7a512c714ae7ecd09eabe347285b25c Mon Sep 17 00:00:00 2001
From: JoostK <joost.koehoorn@gmail.com>
Date: Sat, 8 May 2021 20:29:06 +0200
Subject: [PATCH] fix(platform-browser): prevent memory leak of style nodes if
 shadow DOM encapsulation is used (#42005)

Prior to this change, any inserted `<style>` nodes into shadow dom trees would be retained
in memory, even after the shadow dom tree has been removed. This commit fixes the memory
leak by tracking the inserted `<style>` nodes per host element, such that removal of the
host element also releases the style nodes.

Fixes #36655

PR Close #42005
---
 .../forms_reactive/bundle.golden_symbols.json |  3 ++
 .../bundle.golden_symbols.json                |  3 ++
 .../router/bundle.golden_symbols.json         |  3 ++
 .../src/dom/shared_styles_host.ts             | 30 +++++++++++++------
 .../test/dom/shared_styles_host_spec.ts       |  9 ++++++
 5 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/packages/core/test/bundling/forms_reactive/bundle.golden_symbols.json b/packages/core/test/bundling/forms_reactive/bundle.golden_symbols.json
index b89832be26..ac71b3c288 100644
--- a/packages/core/test/bundling/forms_reactive/bundle.golden_symbols.json
+++ b/packages/core/test/bundling/forms_reactive/bundle.golden_symbols.json
@@ -1430,6 +1430,9 @@
   {
     "name": "removeListItem"
   },
+  {
+    "name": "removeStyle"
+  },
   {
     "name": "renderComponent"
   },
diff --git a/packages/core/test/bundling/forms_template_driven/bundle.golden_symbols.json b/packages/core/test/bundling/forms_template_driven/bundle.golden_symbols.json
index df9a8e4ae4..067827b956 100644
--- a/packages/core/test/bundling/forms_template_driven/bundle.golden_symbols.json
+++ b/packages/core/test/bundling/forms_template_driven/bundle.golden_symbols.json
@@ -1394,6 +1394,9 @@
   {
     "name": "removeListItem"
   },
+  {
+    "name": "removeStyle"
+  },
   {
     "name": "renderComponent"
   },
diff --git a/packages/core/test/bundling/router/bundle.golden_symbols.json b/packages/core/test/bundling/router/bundle.golden_symbols.json
index 2f236a71e8..29c8b54fc2 100644
--- a/packages/core/test/bundling/router/bundle.golden_symbols.json
+++ b/packages/core/test/bundling/router/bundle.golden_symbols.json
@@ -1832,6 +1832,9 @@
   {
     "name": "removeFromArray"
   },
+  {
+    "name": "removeStyle"
+  },
   {
     "name": "renderComponent"
   },
diff --git a/packages/platform-browser/src/dom/shared_styles_host.ts b/packages/platform-browser/src/dom/shared_styles_host.ts
index 25ec7c8816..f82ee229c0 100644
--- a/packages/platform-browser/src/dom/shared_styles_host.ts
+++ b/packages/platform-browser/src/dom/shared_styles_host.ts
@@ -34,35 +34,47 @@ export class SharedStylesHost {
 
 @Injectable()
 export class DomSharedStylesHost extends SharedStylesHost implements OnDestroy {
-  private _hostNodes = new Set<Node>();
-  private _styleNodes = new Set<Node>();
+  // Maps all registered host nodes to a list of style nodes that have been added to the host node.
+  private _hostNodes = new Map<Node, Node[]>();
+
   constructor(@Inject(DOCUMENT) private _doc: any) {
     super();
-    this._hostNodes.add(_doc.head);
+    this._hostNodes.set(_doc.head, []);
   }
 
-  private _addStylesToHost(styles: Set<string>, host: Node): void {
+  private _addStylesToHost(styles: Set<string>, host: Node, styleNodes: Node[]): void {
     styles.forEach((style: string) => {
       const styleEl = this._doc.createElement('style');
       styleEl.textContent = style;
-      this._styleNodes.add(host.appendChild(styleEl));
+      styleNodes.push(host.appendChild(styleEl));
     });
   }
 
   addHost(hostNode: Node): void {
-    this._addStylesToHost(this._stylesSet, hostNode);
-    this._hostNodes.add(hostNode);
+    const styleNodes: Node[] = [];
+    this._addStylesToHost(this._stylesSet, hostNode, styleNodes);
+    this._hostNodes.set(hostNode, styleNodes);
   }
 
   removeHost(hostNode: Node): void {
+    const styleNodes = this._hostNodes.get(hostNode);
+    if (styleNodes) {
+      styleNodes.forEach(removeStyle);
+    }
     this._hostNodes.delete(hostNode);
   }
 
   onStylesAdded(additions: Set<string>): void {
-    this._hostNodes.forEach(hostNode => this._addStylesToHost(additions, hostNode));
+    this._hostNodes.forEach((styleNodes, hostNode) => {
+      this._addStylesToHost(additions, hostNode, styleNodes);
+    });
   }
 
   ngOnDestroy(): void {
-    this._styleNodes.forEach(styleNode => getDOM().remove(styleNode));
+    this._hostNodes.forEach(styleNodes => styleNodes.forEach(removeStyle));
   }
 }
+
+function removeStyle(styleNode: Node): void {
+  getDOM().remove(styleNode);
+}
diff --git a/packages/platform-browser/test/dom/shared_styles_host_spec.ts b/packages/platform-browser/test/dom/shared_styles_host_spec.ts
index b164847f47..ac00e735d1 100644
--- a/packages/platform-browser/test/dom/shared_styles_host_spec.ts
+++ b/packages/platform-browser/test/dom/shared_styles_host_spec.ts
@@ -47,6 +47,15 @@ import {expect} from '@angular/platform-browser/testing/src/matchers';
       expect(doc.head).toHaveText('a {};b {};');
     });
 
+    it('should remove style nodes when the host is removed', () => {
+      ssh.addStyles(['a {};']);
+      ssh.addHost(someHost);
+      expect(someHost.innerHTML).toEqual('<style>a {};</style>');
+
+      ssh.removeHost(someHost);
+      expect(someHost.innerHTML).toEqual('');
+    });
+
     it('should remove style nodes on destroy', () => {
       ssh.addStyles(['a {};']);
       ssh.addHost(someHost);