From 6e736f2452d2e6fdd4da6d8f6f2f44da9d83fa4f Mon Sep 17 00:00:00 2001
From: rinsuki <428rinsuki+git@gmail.com>
Date: Wed, 4 May 2022 10:20:44 +0900
Subject: [PATCH] fix: embed.js doesn't expands iframes height (#18301)

also including some refactoring:
- add `// @ts-check`
- use Map to completely avoid prototype pollution
- assign random id to each iframe for reduce chance to brute-force attack, and leak of iframe counts
- check iframe.contentWindow and MessageEvent.source to validate message is coming from correct iframe (it works on latest Chrome/Firefox/Safari but I'm not sure this is allowed by spec)

follow-up of #17420
fix #18299
---
 public/embed.js | 36 +++++++++++++++++++++++++++++-------
 1 file changed, 29 insertions(+), 7 deletions(-)

diff --git a/public/embed.js b/public/embed.js
index d597fd33c7..5607c24d5f 100644
--- a/public/embed.js
+++ b/public/embed.js
@@ -1,6 +1,11 @@
+// @ts-check
+
 (function() {
   'use strict';
 
+  /**
+   * @param {() => void} loaded
+   */
   var ready = function(loaded) {
     if (['interactive', 'complete'].indexOf(document.readyState) !== -1) {
       loaded();
@@ -10,26 +15,43 @@
   };
 
   ready(function() {
-    var iframes = [];
+    /** @type {Map<number, HTMLIFrameElement>} */
+    var iframes = new Map();
 
     window.addEventListener('message', function(e) {
       var data = e.data || {};
 
-      if (data.type !== 'setHeight' || !iframes[data.id] || window.location.origin !== e.origin || data.id.toString() === '__proto__') {
+      if (typeof data !== 'object' || data.type !== 'setHeight' || !iframes.has(data.id)) {
         return;
       }
 
-      iframes[data.id].height = data.height;
+      var iframe = iframes.get(data.id);
+
+      if ('source' in e && iframe.contentWindow !== e.source) {
+        return;
+      }
+
+      iframe.height = data.height;
     });
 
     [].forEach.call(document.querySelectorAll('iframe.mastodon-embed'), function(iframe) {
+      // select unique id for each iframe
+      var id = 0, failCount = 0, idBuffer = new Uint32Array(1);
+      while (id === 0 || iframes.has(id)) {
+        id = crypto.getRandomValues(idBuffer)[0];
+        failCount++;
+        if (failCount > 100) {
+          // give up and assign (easily guessable) unique number if getRandomValues is broken or no luck
+          id = -(iframes.size + 1);
+          break;
+        }
+      }
+
+      iframes.set(id, iframe);
+
       iframe.scrolling      = 'no';
       iframe.style.overflow = 'hidden';
 
-      iframes.push(iframe);
-
-      var id = iframes.length - 1;
-
       iframe.onload = function() {
         iframe.contentWindow.postMessage({
           type: 'setHeight',