From 1eebb57bf41c2e3907c9456c1866e8bcc7110b9d Mon Sep 17 00:00:00 2001
From: K <z@kwi.li>
Date: Tue, 28 Dec 2021 15:49:17 -0500
Subject: [PATCH] Add RequireSandboxOnIFrame

---
 policy.go        | 72 ++++++++++++++++++++++++++++++++++++++++++++++++
 sanitize.go      | 25 ++++++++++++++++-
 sanitize_test.go | 18 ++++++++++++
 3 files changed, 114 insertions(+), 1 deletion(-)

diff --git a/policy.go b/policy.go
index 71f6b8d..c412776 100644
--- a/policy.go
+++ b/policy.go
@@ -74,6 +74,9 @@ type Policy struct {
 	// When true, add crossorigin="anonymous" to HTML audio, img, link, script, and video tags
 	requireCrossOriginAnonymous bool
 
+	// When true, add and filter sandbox attribute on iframe tags
+	requireSandboxOnIFrame map[string]bool
+
 	// When true add target="_blank" to fully qualified links
 	// Will add for href="http://foo"
 	// Will skip for href="/foo" or href="foo"
@@ -189,6 +192,25 @@ type stylePolicyBuilder struct {
 
 type urlPolicy func(url *url.URL) (allowUrl bool)
 
+type SandboxValue int64
+
+const (
+	SandboxAllowDownloads SandboxValue = iota
+	SandboxAllowDownloadsWithoutUserActivation
+	SandboxAllowForms
+	SandboxAllowModals
+	SandboxAllowOrientationLock
+	SandboxAllowPointerLock
+	SandboxAllowPopups
+	SandboxAllowPopupsToEscapeSandbox
+	SandboxAllowPresentation
+	SandboxAllowSameOrigin
+	SandboxAllowScripts
+	SandboxAllowStorageAccessByUserActivation
+	SandboxAllowTopNavigation
+	SandboxAllowTopNavigationByUserActivation
+)
+
 // init initializes the maps if this has not been done already
 func (p *Policy) init() {
 	if !p.initialized {
@@ -680,6 +702,56 @@ func (p *Policy) AllowURLSchemeWithCustomPolicy(
 	return p
 }
 
+func (p *Policy) RequireSandboxOnIFrame(vals ...SandboxValue) {
+	p.requireSandboxOnIFrame = make(map[string]bool)
+
+	for val := range vals {
+		switch SandboxValue(val) {
+		case SandboxAllowDownloads:
+			p.requireSandboxOnIFrame["allow-downloads"] = true
+
+		case SandboxAllowDownloadsWithoutUserActivation:
+			p.requireSandboxOnIFrame["allow-downloads-without-user-activation"] = true
+
+		case SandboxAllowForms:
+			p.requireSandboxOnIFrame["allow-forms"] = true
+
+		case SandboxAllowModals:
+			p.requireSandboxOnIFrame["allow-modals"] = true
+
+		case SandboxAllowOrientationLock:
+			p.requireSandboxOnIFrame["allow-orientation-lock"] = true
+
+		case SandboxAllowPointerLock:
+			p.requireSandboxOnIFrame["allow-pointer-lock"] = true
+
+		case SandboxAllowPopups:
+			p.requireSandboxOnIFrame["allow-popups"] = true
+
+		case SandboxAllowPopupsToEscapeSandbox:
+			p.requireSandboxOnIFrame["allow-popups-to-escape-sandbox"] = true
+
+		case SandboxAllowPresentation:
+			p.requireSandboxOnIFrame["allow-presentation"] = true
+
+		case SandboxAllowSameOrigin:
+			p.requireSandboxOnIFrame["allow-same-origin"] = true
+
+		case SandboxAllowScripts:
+			p.requireSandboxOnIFrame["allow-scripts"] = true
+
+		case SandboxAllowStorageAccessByUserActivation:
+			p.requireSandboxOnIFrame["allow-storage-access-by-user-activation"] = true
+
+		case SandboxAllowTopNavigation:
+			p.requireSandboxOnIFrame["allow-top-navigation"] = true
+
+		case SandboxAllowTopNavigationByUserActivation:
+			p.requireSandboxOnIFrame["allow-top-navigation-by-user-activation"] = true
+		}
+	}
+}
+
 // AddSpaceWhenStrippingTag states whether to add a single space " " when
 // removing tags that are not allowed by the policy.
 //
diff --git a/sanitize.go b/sanitize.go
index 97628ce..b888625 100644
--- a/sanitize.go
+++ b/sanitize.go
@@ -240,7 +240,7 @@ func (p *Policy) sanitize(r io.Reader, w io.Writer) error {
 	// rather than:
 	//   p := bluemonday.NewPolicy()
 	// If this is the case, and if they haven't yet triggered an action that
-	// would initiliaze the maps, then we need to do that.
+	// would initialize the maps, then we need to do that.
 	p.init()
 
 	buff, ok := w.(stringWriterWriter)
@@ -809,6 +809,29 @@ attrsLoop:
 		}
 	}
 
+	if p.requireSandboxOnIFrame != nil && elementName == "iframe" {
+		var sandboxFound bool
+		for i, htmlAttr := range cleanAttrs {
+			if htmlAttr.Key == "sandbox" {
+				sandboxFound = true
+				var cleanVals []string
+				for _, val := range strings.Fields(htmlAttr.Val) {
+					if p.requireSandboxOnIFrame[val] {
+						cleanVals = append(cleanVals, val)
+					}
+				}
+				cleanAttrs[i].Val = strings.Join(cleanVals, " ")
+			}
+		}
+
+		if !sandboxFound {
+			sandbox := html.Attribute{}
+			sandbox.Key = "sandbox"
+			sandbox.Val = ""
+			cleanAttrs = append(cleanAttrs, sandbox)
+		}
+	}
+
 	return cleanAttrs
 }
 
diff --git a/sanitize_test.go b/sanitize_test.go
index 0905bd8..c322ae5 100644
--- a/sanitize_test.go
+++ b/sanitize_test.go
@@ -1871,6 +1871,24 @@ func TestIssue107(t *testing.T) {
 	wg.Wait()
 }
 
+func TestIFrameSandbox(t *testing.T) {
+	p := NewPolicy()
+	p.AllowAttrs("sandbox").OnElements("iframe")
+	p.RequireSandboxOnIFrame(SandboxAllowDownloads)
+
+	in := `<iframe src="http://example.com" sandbox="allow-forms allow-downloads"></iframe>`
+	expected := `<iframe sandbox="allow-downloads"></iframe>`
+	out := p.Sanitize(in)
+	if out != expected {
+		t.Errorf(
+			"test failed;\ninput   : %s\noutput  : %s\nexpected: %s",
+			in,
+			out,
+			expected,
+		)
+	}
+}
+
 func TestSanitizedURL(t *testing.T) {
 	tests := []test{
 		{