google · copybara-service · Oct 1, 2024
diff --git a/src/dom/elements/script.ts b/src/dom/elements/script.ts
@@ -13,8 +13,7 @@ import {getScriptNonce} from '../globals/window.js';
 
 /** Propagates CSP nonce to dynamically created scripts. */
 function setNonceForScriptElement(script: HTMLScriptElement) {
-  const win = script.ownerDocument && script.ownerDocument.defaultView;
-  const nonce = getScriptNonce(win || window);
+  const nonce = getScriptNonce(script.ownerDocument || document);
   if (nonce) {
     script.setAttribute('nonce', nonce);
   }

diff --git a/src/dom/globals/window.ts b/src/dom/globals/window.ts
@@ -23,28 +23,43 @@ export function open(
   return null;
 }
 
-/** Returns CSP nonce, if set for any script tag. */
-export function getScriptNonce(win: Window): string {
-  return getNonceFor('script', win);
+/**
+ * Returns CSP nonce, if set for any script tag.
+ */
+export function getScriptNonce(documentOrWindow?: Document | Window): string {
+  return getNonceFor('script', documentOrWindow);
 }
 
-/** Returns CSP nonce, if set for any style tag. */
-export function getStyleNonce(win: Window): string {
-  return getNonceFor('style', win);
+/**
+ * Returns CSP nonce, if set for any style tag.
+ */
+export function getStyleNonce(documentOrWindow?: Document | Window): string {
+  return getNonceFor('style', documentOrWindow);
 }
 
-function getNonceFor(elementName: 'script' | 'style', win: Window): string {
-  const doc = win.document;
+function getNonceFor(
+  elementName: 'script' | 'style',
+  documentOrWindow: Document | Window = document,
+): string {
+  // Note: We can't use `instanceof` here because `documentOrWindow` likely
+  // comes from a different realm.
+  const doc =
+    'document' in documentOrWindow
+      ? documentOrWindow.document
+      : documentOrWindow;
+
   // document.querySelector can be undefined in non-browser environments.
   const el = doc.querySelector?.<HTMLScriptElement | HTMLStyleElement>(
     `${elementName}[nonce]`,
   );
-  if (el) {
-    // Try to get the nonce from the IDL property first, because browsers that
-    // implement additional nonce protection features (currently only Chrome) to
-    // prevent nonce stealing via CSS do not expose the nonce via attributes.
-    // See https://github.com/whatwg/html/issues/2369
-    return el['nonce'] || el.getAttribute('nonce') || '';
+
+  if (el == null) {
+    return '';
   }
-  return '';
+
+  // Try to get the nonce from the IDL property first, because browsers that
+  // implement additional nonce protection features (currently only Chrome) to
+  // prevent nonce stealing via CSS do not expose the nonce via attributes.
+  // See https://github.com/whatwg/html/issues/2369
+  return el['nonce'] || el.getAttribute('nonce') || '';
 }
diff --git a/test/dom/globals/window_test.ts b/test/dom/globals/window_test.ts
@@ -0,0 +1,60 @@
+/**
+ * @license
+ * Copyright Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {getScriptNonce} from '../../../src/dom/globals/window';
+
+describe('Window', () => {
+  describe('getScriptNonce', () => {
+    it('returns a nonce if a script tag with nonce is found', () => {
+      const doc = document.implementation.createHTMLDocument();
+      const script = doc.createElement('script');
+      script.setAttribute('nonce', '123');
+      doc.body.appendChild(script);
+      expect(getScriptNonce(doc)).toEqual('123');
+    });
+
+    it('accepts a window', async () => {
+      const iframe = document.createElement('iframe');
+      const iframeLoaded = new Promise((resolve) => {
+        iframe.>
+      });
+      document.body.appendChild(iframe);
+      await iframeLoaded;
+
+      const doc = iframe.contentDocument!;
+      const script = doc.createElement('script');
+      script.setAttribute('nonce', '234');
+      doc.body.appendChild(script);
+
+      const win = iframe.contentWindow!;
+      expect(win).not.toBeNull();
+      expect(getScriptNonce(win)).toEqual('234');
+    });
+
+    it('returns empty string if no script tag is found', () => {
+      const doc = document.implementation.createHTMLDocument();
+      expect(getScriptNonce(doc)).toEqual('');
+    });
+
+    it('returns empty string if no nonce is found', () => {
+      const doc = document.implementation.createHTMLDocument();
+      const script = doc.createElement('script');
+      doc.body.appendChild(script);
+      expect(getScriptNonce(doc)).toEqual('');
+    });
+
+    it('returns nonce from current document when passed in no arguments', () => {
+      const script = document.createElement('script');
+      script.setAttribute('nonce', '345');
+
+      spyOn(document, 'querySelector')
+        .withArgs('script[nonce]')
+        .and.returnValue(script);
+
+      expect(getScriptNonce()).toEqual('345');
+    });
+  });
+});