Repository: microcosm-cc/bluemonday
Pull request #125 (2021-06-16 19:37:06)

Additive policies
This PR enables additive policies with multiple policies per element and attribute.

Instead of
`p.AllowAttrs("class").Matching(regexp.MustCompile("red|green")).OnElements("span")`
you can write now
```
p.AllowAttrs("class").Matching(regexp.MustCompile("red")).OnElements("span")
p.AllowAttrs("class").Matching(regexp.MustCompile("green")).OnElements("span")
```
This looks worse at the moment but an example use case are different renderers which add their rules to the policy. Think of an interface method `AddRules(p *Policy)` and Renderer1 wants to allow the class `red` and Renderer2 wants to allow the class `green`. Both call `p.AllowAttrs("class").Matching(regexp.MustCompile("xxx")).OnElements("span")` but only the last renderer wins. With this PR the sanitizer allows a attribute if one of the rules allows it.

Or [this error](https://github.com/go-gitea/gitea/pull/16163)


Patch:
diff --git a/policy.go b/policy.go
index df6152d..6a942c8 100644
--- a/policy.go
+++ b/policy.go
@@ -91,28 +91,28 @@ type Policy struct {
 	// When true, allow comments.
 	allowComments bool
 
-	// map[htmlElementName]map[htmlAttributeName]attrPolicy
-	elsAndAttrs map[string]map[string]attrPolicy
+	// map[htmlElementName]map[htmlAttributeName][]attrPolicy
+	elsAndAttrs map[string]map[string][]attrPolicy
 
 	// elsMatchingAndAttrs stores regex based element matches along with attributes
-	elsMatchingAndAttrs map[*regexp.Regexp]map[string]attrPolicy
+	elsMatchingAndAttrs map[*regexp.Regexp]map[string][]attrPolicy
 
-	// map[htmlAttributeName]attrPolicy
-	globalAttrs map[string]attrPolicy
+	// map[htmlAttributeName][]attrPolicy
+	globalAttrs map[string][]attrPolicy
 
-	// map[htmlElementName]map[cssPropertyName]stylePolicy
-	elsAndStyles map[string]map[string]stylePolicy
+	// map[htmlElementName]map[cssPropertyName][]stylePolicy
+	elsAndStyles map[string]map[string][]stylePolicy
 
-	// map[regex]map[cssPropertyName]stylePolicy
-	elsMatchingAndStyles map[*regexp.Regexp]map[string]stylePolicy
+	// map[regex]map[cssPropertyName][]stylePolicy
+	elsMatchingAndStyles map[*regexp.Regexp]map[string][]stylePolicy
 
-	// map[cssPropertyName]stylePolicy
-	globalStyles map[string]stylePolicy
+	// map[cssPropertyName][]stylePolicy
+	globalStyles map[string][]stylePolicy
 
 	// If urlPolicy is nil, all URLs with matching schema are allowed.
 	// Otherwise, only the URLs with matching schema and urlPolicy(url)
 	// returning true are allowed.
-	allowURLSchemes map[string]urlPolicy
+	allowURLSchemes map[string][]urlPolicy
 
 	// If an element has had all attributes removed as a result of a policy
 	// being applied, then the element would be removed from the output.
@@ -179,13 +179,13 @@ type urlPolicy func(url *url.URL) (allowUrl bool)
 // init initializes the maps if this has not been done already
 func (p *Policy) init() {
 	if !p.initialized {
-		p.elsAndAttrs = make(map[string]map[string]attrPolicy)
-		p.elsMatchingAndAttrs = make(map[*regexp.Regexp]map[string]attrPolicy)
-		p.globalAttrs = make(map[string]attrPolicy)
-		p.elsAndStyles = make(map[string]map[string]stylePolicy)
-		p.elsMatchingAndStyles = make(map[*regexp.Regexp]map[string]stylePolicy)
-		p.globalStyles = make(map[string]stylePolicy)
-		p.allowURLSchemes = make(map[string]urlPolicy)
+		p.elsAndAttrs = make(map[string]map[string][]attrPolicy)
+		p.elsMatchingAndAttrs = make(map[*regexp.Regexp]map[string][]attrPolicy)
+		p.globalAttrs = make(map[string][]attrPolicy)
+		p.elsAndStyles = make(map[string]map[string][]stylePolicy)
+		p.elsMatchingAndStyles = make(map[*regexp.Regexp]map[string][]stylePolicy)
+		p.globalStyles = make(map[string][]stylePolicy)
+		p.allowURLSchemes = make(map[string][]urlPolicy)
 		p.setOfElementsAllowedWithoutAttrs = make(map[string]struct{})
 		p.setOfElementsToSkipContent = make(map[string]struct{})
 		p.initialized = true
@@ -286,8 +286,7 @@ func (abp *attrPolicyBuilder) AllowNoAttrs() *attrPolicyBuilder {
 }
 
 // Matching allows a regular expression to be applied to a nascent attribute
-// policy, and returns the attribute policy. Calling this more than once will
-// replace the existing regexp.
+// policy, and returns the attribute policy.
 func (abp *attrPolicyBuilder) Matching(regex *regexp.Regexp) *attrPolicyBuilder {
 
 	abp.regexp = regex
@@ -305,7 +304,7 @@ func (abp *attrPolicyBuilder) OnElements(elements ...string) *Policy {
 		for _, attr := range abp.attrNames {
 
 			if _, ok := abp.p.elsAndAttrs[element]; !ok {
-				abp.p.elsAndAttrs[element] = make(map[string]attrPolicy)
+				abp.p.elsAndAttrs[element] = make(map[string][]attrPolicy)
 			}
 
 			ap := attrPolicy{}
@@ -313,14 +312,14 @@ func (abp *attrPolicyBuilder) OnElements(elements ...string) *Policy {
 				ap.regexp = abp.regexp
 			}
 
-			abp.p.elsAndAttrs[element][attr] = ap
+			abp.p.elsAndAttrs[element][attr] = append(abp.p.elsAndAttrs[element][attr], ap)
 		}
 
 		if abp.allowEmpty {
 			abp.p.setOfElementsAllowedWithoutAttrs[element] = struct{}{}
 
 			if _, ok := abp.p.elsAndAttrs[element]; !ok {
-				abp.p.elsAndAttrs[element] = make(map[string]attrPolicy)
+				abp.p.elsAndAttrs[element] = make(map[string][]attrPolicy)
 			}
 		}
 	}
@@ -333,19 +332,19 @@ func (abp *attrPolicyBuilder) OnElements(elements ...string) *Policy {
 func (abp *attrPolicyBuilder) OnElementsMatching(regex *regexp.Regexp) *Policy {
 	for _, attr := range abp.attrNames {
 		if _, ok := abp.p.elsMatchingAndAttrs[regex]; !ok {
-			abp.p.elsMatchingAndAttrs[regex] = make(map[string]attrPolicy)
+			abp.p.elsMatchingAndAttrs[regex] = make(map[string][]attrPolicy)
 		}
 		ap := attrPolicy{}
 		if abp.regexp != nil {
 			ap.regexp = abp.regexp
 		}
-		abp.p.elsMatchingAndAttrs[regex][attr] = ap
+		abp.p.elsMatchingAndAttrs[regex][attr] = append(abp.p.elsMatchingAndAttrs[regex][attr], ap)
 	}
 
 	if abp.allowEmpty {
 		abp.p.setOfElementsMatchingAllowedWithoutAttrs = append(abp.p.setOfElementsMatchingAllowedWithoutAttrs, regex)
 		if _, ok := abp.p.elsMatchingAndAttrs[regex]; !ok {
-			abp.p.elsMatchingAndAttrs[regex] = make(map[string]attrPolicy)
+			abp.p.elsMatchingAndAttrs[regex] = make(map[string][]attrPolicy)
 		}
 	}
 
@@ -358,7 +357,7 @@ func (abp *attrPolicyBuilder) Globally() *Policy {
 
 	for _, attr := range abp.attrNames {
 		if _, ok := abp.p.globalAttrs[attr]; !ok {
-			abp.p.globalAttrs[attr] = attrPolicy{}
+			abp.p.globalAttrs[attr] = []attrPolicy{}
 		}
 
 		ap := attrPolicy{}
@@ -366,7 +365,7 @@ func (abp *attrPolicyBuilder) Globally() *Policy {
 			ap.regexp = abp.regexp
 		}
 
-		abp.p.globalAttrs[attr] = ap
+		abp.p.globalAttrs[attr] = append(abp.p.globalAttrs[attr], ap)
 	}
 
 	return abp.p
@@ -394,8 +393,7 @@ func (p *Policy) AllowStyles(propertyNames ...string) *stylePolicyBuilder {
 }
 
 // Matching allows a regular expression to be applied to a nascent style
-// policy, and returns the style policy. Calling this more than once will
-// replace the existing regexp.
+// policy, and returns the style policy.
 func (spb *stylePolicyBuilder) Matching(regex *regexp.Regexp) *stylePolicyBuilder {
 
 	spb.regexp = regex
@@ -404,8 +402,7 @@ func (spb *stylePolicyBuilder) Matching(regex *regexp.Regexp) *stylePolicyBuilde
 }
 
 // MatchingEnum allows a list of allowed values to be applied to a nascent style
-// policy, and returns the style policy. Calling this more than once will
-// replace the existing list of allowed values.
+// policy, and returns the style policy.
 func (spb *stylePolicyBuilder) MatchingEnum(enum ...string) *stylePolicyBuilder {
 
 	spb.enum = enum
@@ -414,8 +411,7 @@ func (spb *stylePolicyBuilder) MatchingEnum(enum ...string) *stylePolicyBuilder
 }
 
 // MatchingHandler allows a handler to be applied to a nascent style
-// policy, and returns the style policy. Calling this more than once will
-// replace the existing handler.
+// policy, and returns the style policy.
 func (spb *stylePolicyBuilder) MatchingHandler(handler func(string) bool) *stylePolicyBuilder {
 
 	spb.handler = handler
@@ -433,7 +429,7 @@ func (spb *stylePolicyBuilder) OnElements(elements ...string) *Policy {
 		for _, attr := range spb.propertyNames {
 
 			if _, ok := spb.p.elsAndStyles[element]; !ok {
-				spb.p.elsAndStyles[element] = make(map[string]stylePolicy)
+				spb.p.elsAndStyles[element] = make(map[string][]stylePolicy)
 			}
 
 			sp := stylePolicy{}
@@ -446,7 +442,7 @@ func (spb *stylePolicyBuilder) OnElements(elements ...string) *Policy {
 			} else {
 				sp.handler = css.GetDefaultHandler(attr)
 			}
-			spb.p.elsAndStyles[element][attr] = sp
+			spb.p.elsAndStyles[element][attr] = append(spb.p.elsAndStyles[element][attr], sp)
 		}
 	}
 
@@ -460,7 +456,7 @@ func (spb *stylePolicyBuilder) OnElementsMatching(regex *regexp.Regexp) *Policy
 	for _, attr := range spb.propertyNames {
 
 		if _, ok := spb.p.elsMatchingAndStyles[regex]; !ok {
-			spb.p.elsMatchingAndStyles[regex] = make(map[string]stylePolicy)
+			spb.p.elsMatchingAndStyles[regex] = make(map[string][]stylePolicy)
 		}
 
 		sp := stylePolicy{}
@@ -473,7 +469,7 @@ func (spb *stylePolicyBuilder) OnElementsMatching(regex *regexp.Regexp) *Policy
 		} else {
 			sp.handler = css.GetDefaultHandler(attr)
 		}
-		spb.p.elsMatchingAndStyles[regex][attr] = sp
+		spb.p.elsMatchingAndStyles[regex][attr] = append(spb.p.elsMatchingAndStyles[regex][attr], sp)
 	}
 
 	return spb.p
@@ -485,7 +481,7 @@ func (spb *stylePolicyBuilder) Globally() *Policy {
 
 	for _, attr := range spb.propertyNames {
 		if _, ok := spb.p.globalStyles[attr]; !ok {
-			spb.p.globalStyles[attr] = stylePolicy{}
+			spb.p.globalStyles[attr] = []stylePolicy{}
 		}
 
 		// Use only one strategy for validating styles, fallback to default
@@ -499,7 +495,7 @@ func (spb *stylePolicyBuilder) Globally() *Policy {
 		} else {
 			sp.handler = css.GetDefaultHandler(attr)
 		}
-		spb.p.globalStyles[attr] = sp
+		spb.p.globalStyles[attr] = append(spb.p.globalStyles[attr], sp)
 	}
 
 	return spb.p
@@ -515,7 +511,7 @@ func (p *Policy) AllowElements(names ...string) *Policy {
 		element = strings.ToLower(element)
 
 		if _, ok := p.elsAndAttrs[element]; !ok {
-			p.elsAndAttrs[element] = make(map[string]attrPolicy)
+			p.elsAndAttrs[element] = make(map[string][]attrPolicy)
 		}
 	}
 
@@ -525,7 +521,7 @@ func (p *Policy) AllowElements(names ...string) *Policy {
 func (p *Policy) AllowElementsMatching(regex *regexp.Regexp) *Policy {
 	p.init()
 	if _, ok := p.elsMatchingAndAttrs[regex]; !ok {
-		p.elsMatchingAndAttrs[regex] = make(map[string]attrPolicy)
+		p.elsMatchingAndAttrs[regex] = make(map[string][]attrPolicy)
 	}
 	return p
 }
@@ -664,7 +660,7 @@ func (p *Policy) AllowURLSchemeWithCustomPolicy(
 
 	scheme = strings.ToLower(scheme)
 
-	p.allowURLSchemes[scheme] = urlPolicy
+	p.allowURLSchemes[scheme] = append(p.allowURLSchemes[scheme], urlPolicy)
 
 	return p
 }
diff --git a/sanitize.go b/sanitize.go
index 1d85485..eee5db6 100644
--- a/sanitize.go
+++ b/sanitize.go
@@ -445,7 +445,7 @@ func (p *Policy) sanitize(r io.Reader) *bytes.Buffer {
 func (p *Policy) sanitizeAttrs(
 	elementName string,
 	attrs []html.Attribute,
-	aps map[string]attrPolicy,
+	aps map[string][]attrPolicy,
 ) []html.Attribute {
 
 	if len(attrs) == 0 {
@@ -472,6 +472,7 @@ func (p *Policy) sanitizeAttrs(
 	// Builds a new attribute slice based on the whether the attribute has been
 	// whitelisted explicitly or globally.
 	cleanAttrs := []html.Attribute{}
+attrsLoop:
 	for _, htmlAttr := range attrs {
 		if p.allowDataAttributes {
 			// If we see a data attribute, let it through.
@@ -494,27 +495,30 @@ func (p *Policy) sanitizeAttrs(
 		}
 
 		// Is there an element specific attribute policy that applies?
-		if ap, ok := aps[htmlAttr.Key]; ok {
-			if ap.regexp != nil {
-				if ap.regexp.MatchString(htmlAttr.Val) {
+		if apl, ok := aps[htmlAttr.Key]; ok {
+			for _, ap := range apl {
+				if ap.regexp != nil {
+					if ap.regexp.MatchString(htmlAttr.Val) {
+						cleanAttrs = append(cleanAttrs, htmlAttr)
+						continue attrsLoop
+					}
+				} else {
 					cleanAttrs = append(cleanAttrs, htmlAttr)
-					continue
+					continue attrsLoop
 				}
-			} else {
-				cleanAttrs = append(cleanAttrs, htmlAttr)
-				continue
 			}
 		}
 
 		// Is there a global attribute policy that applies?
-		if ap, ok := p.globalAttrs[htmlAttr.Key]; ok {
-
-			if ap.regexp != nil {
-				if ap.regexp.MatchString(htmlAttr.Val) {
+		if apl, ok := p.globalAttrs[htmlAttr.Key]; ok {
+			for _, ap := range apl {
+				if ap.regexp != nil {
+					if ap.regexp.MatchString(htmlAttr.Val) {
+						cleanAttrs = append(cleanAttrs, htmlAttr)
+					}
+				} else {
 					cleanAttrs = append(cleanAttrs, htmlAttr)
 				}
-			} else {
-				cleanAttrs = append(cleanAttrs, htmlAttr)
 			}
 		}
 	}
@@ -758,14 +762,14 @@ func (p *Policy) sanitizeAttrs(
 func (p *Policy) sanitizeStyles(attr html.Attribute, elementName string) html.Attribute {
 	sps := p.elsAndStyles[elementName]
 	if len(sps) == 0 {
-		sps = map[string]stylePolicy{}
+		sps = map[string][]stylePolicy{}
 		// check for any matching elements, if we don't already have a policy found
 		// if multiple matches are found they will be overwritten, it's best
 		// to not have overlapping matchers
 		for regex, policies := range p.elsMatchingAndStyles {
 			if regex.MatchString(elementName) {
 				for k, v := range policies {
-					sps[k] = v
+					sps[k] = append(sps[k], v...)
 				}
 			}
 		}
@@ -783,46 +787,51 @@ func (p *Policy) sanitizeStyles(attr html.Attribute, elementName string) html.At
 	clean := []string{}
 	prefixes := []string{"-webkit-", "-moz-", "-ms-", "-o-", "mso-", "-xv-", "-atsc-", "-wap-", "-khtml-", "prince-", "-ah-", "-hp-", "-ro-", "-rim-", "-tc-"}
 
+decLoop:
 	for _, dec := range decs {
-		addedProperty := false
 		tempProperty := strings.ToLower(dec.Property)
 		tempValue := removeUnicode(strings.ToLower(dec.Value))
 		for _, i := range prefixes {
 			tempProperty = strings.TrimPrefix(tempProperty, i)
 		}
-		if sp, ok := sps[tempProperty]; ok {
-			if sp.handler != nil {
-				if sp.handler(tempValue) {
-					clean = append(clean, dec.Property+": "+dec.Value)
-					addedProperty = true
-				}
-			} else if len(sp.enum) > 0 {
-				if stringInSlice(tempValue, sp.enum) {
-					clean = append(clean, dec.Property+": "+dec.Value)
-					addedProperty = true
-				}
-			} else if sp.regexp != nil {
-				if sp.regexp.MatchString(tempValue) {
-					clean = append(clean, dec.Property+": "+dec.Value)
-					addedProperty = true
+		if spl, ok := sps[tempProperty]; ok {
+			for _, sp := range spl {
+				if sp.handler != nil {
+					if sp.handler(tempValue) {
+						clean = append(clean, dec.Property+": "+dec.Value)
+						continue decLoop
+					}
+				} else if len(sp.enum) > 0 {
+					if stringInSlice(tempValue, sp.enum) {
+						clean = append(clean, dec.Property+": "+dec.Value)
+						continue decLoop
+					}
+				} else if sp.regexp != nil {
+					if sp.regexp.MatchString(tempValue) {
+						clean = append(clean, dec.Property+": "+dec.Value)
+						continue decLoop
+					}
 				}
-				continue
 			}
 		}
-		if sp, ok := p.globalStyles[tempProperty]; ok && !addedProperty {
-			if sp.handler != nil {
-				if sp.handler(tempValue) {
-					clean = append(clean, dec.Property+": "+dec.Value)
-				}
-			} else if len(sp.enum) > 0 {
-				if stringInSlice(tempValue, sp.enum) {
-					clean = append(clean, dec.Property+": "+dec.Value)
-				}
-			} else if sp.regexp != nil {
-				if sp.regexp.MatchString(tempValue) {
-					clean = append(clean, dec.Property+": "+dec.Value)
+		if spl, ok := p.globalStyles[tempProperty]; ok {
+			for _, sp := range spl {
+				if sp.handler != nil {
+					if sp.handler(tempValue) {
+						clean = append(clean, dec.Property+": "+dec.Value)
+						continue decLoop
+					}
+				} else if len(sp.enum) > 0 {
+					if stringInSlice(tempValue, sp.enum) {
+						clean = append(clean, dec.Property+": "+dec.Value)
+						continue decLoop
+					}
+				} else if sp.regexp != nil {
+					if sp.regexp.MatchString(tempValue) {
+						clean = append(clean, dec.Property+": "+dec.Value)
+						continue decLoop
+					}
 				}
-				continue
 			}
 		}
 	}
@@ -885,15 +894,21 @@ func (p *Policy) validURL(rawurl string) (string, bool) {
 
 		if u.Scheme != "" {
 
-			urlPolicy, ok := p.allowURLSchemes[u.Scheme]
+			urlPolicies, ok := p.allowURLSchemes[u.Scheme]
 			if !ok {
 				return "", false
 			}
 
-			if urlPolicy == nil || urlPolicy(u) == true {
+			if len(urlPolicies) == 0 {
 				return u.String(), true
 			}
 
+			for _, urlPolicy := range urlPolicies {
+				if urlPolicy(u) == true {
+					return u.String(), true
+				}
+			}
+
 			return "", false
 		}
 
@@ -985,14 +1000,14 @@ func removeUnicode(value string) string {
 	return substitutedValue
 }
 
-func (p *Policy) matchRegex(elementName string) (map[string]attrPolicy, bool) {
-	aps := make(map[string]attrPolicy, 0)
+func (p *Policy) matchRegex(elementName string) (map[string][]attrPolicy, bool) {
+	aps := make(map[string][]attrPolicy, 0)
 	matched := false
 	for regex, attrs := range p.elsMatchingAndAttrs {
 		if regex.MatchString(elementName) {
 			matched = true
 			for k, v := range attrs {
-				aps[k] = v
+				aps[k] = append(aps[k], v...)
 			}
 		}
 	}


Test patch:
diff --git a/sanitize_test.go b/sanitize_test.go
index 01260d3..1682756 100644
--- a/sanitize_test.go
+++ b/sanitize_test.go
@@ -3307,3 +3307,320 @@ func TestStyleBlockHandler(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+func TestAdditivePolicies(t *testing.T) {
+	t.Run("AllowAttrs", func(t *testing.T) {
+		p := NewPolicy()
+		p.AllowAttrs("class").Matching(regexp.MustCompile("red")).OnElements("span")
+
+		t.Run("red", func(t *testing.T) {
+			tests := []test{
+				{
+					in:       `<span class="red">test</span>`,
+					expected: `<span class="red">test</span>`,
+				},
+				{
+					in:       `<span class="green">test</span>`,
+					expected: `<span>test</span>`,
+				},
+				{
+					in:       `<span class="blue">test</span>`,
+					expected: `<span>test</span>`,
+				},
+			}
+
+			for ii, tt := range tests {
+				out := p.Sanitize(tt.in)
+				if out != tt.expected {
+					t.Errorf(
+						"test %d failed;\ninput   : %s\noutput  : %s\nexpected: %s",
+						ii,
+						tt.in,
+						out,
+						tt.expected,
+					)
+				}
+			}
+		})
+
+		p.AllowAttrs("class").Matching(regexp.MustCompile("green")).OnElements("span")
+
+		t.Run("green", func(t *testing.T) {
+			tests := []test{
+				{
+					in:       `<span class="red">test</span>`,
+					expected: `<span class="red">test</span>`,
+				},
+				{
+					in:       `<span class="green">test</span>`,
+					expected: `<span class="green">test</span>`,
+				},
+				{
+					in:       `<span class="blue">test</span>`,
+					expected: `<span>test</span>`,
+				},
+			}
+
+			for ii, tt := range tests {
+				out := p.Sanitize(tt.in)
+				if out != tt.expected {
+					t.Errorf(
+						"test %d failed;\ninput   : %s\noutput  : %s\nexpected: %s",
+						ii,
+						tt.in,
+						out,
+						tt.expected,
+					)
+				}
+			}
+		})
+
+		p.AllowAttrs("class").Matching(regexp.MustCompile("yellow")).OnElements("span")
+
+		t.Run("yellow", func(t *testing.T) {
+			tests := []test{
+				{
+					in:       `<span class="red">test</span>`,
+					expected: `<span class="red">test</span>`,
+				},
+				{
+					in:       `<span class="green">test</span>`,
+					expected: `<span class="green">test</span>`,
+				},
+				{
+					in:       `<span class="blue">test</span>`,
+					expected: `<span>test</span>`,
+				},
+			}
+
+			for ii, tt := range tests {
+				out := p.Sanitize(tt.in)
+				if out != tt.expected {
+					t.Errorf(
+						"test %d failed;\ninput   : %s\noutput  : %s\nexpected: %s",
+						ii,
+						tt.in,
+						out,
+						tt.expected,
+					)
+				}
+			}
+		})
+	})
+
+	t.Run("AllowStyles", func(t *testing.T) {
+		p := NewPolicy()
+		p.AllowAttrs("style").OnElements("span")
+		p.AllowStyles("color").Matching(regexp.MustCompile("red")).OnElements("span")
+
+		t.Run("red", func(t *testing.T) {
+			tests := []test{
+				{
+					in:       `<span style="color: red">test</span>`,
+					expected: `<span style="color: red">test</span>`,
+				},
+				{
+					in:       `<span style="color: green">test</span>`,
+					expected: `<span>test</span>`,
+				},
+				{
+					in:       `<span style="color: blue">test</span>`,
+					expected: `<span>test</span>`,
+				},
+			}
+
+			for ii, tt := range tests {
+				out := p.Sanitize(tt.in)
+				if out != tt.expected {
+					t.Errorf(
+						"test %d failed;\ninput   : %s\noutput  : %s\nexpected: %s",
+						ii,
+						tt.in,
+						out,
+						tt.expected,
+					)
+				}
+			}
+		})
+
+		p.AllowStyles("color").Matching(regexp.MustCompile("green")).OnElements("span")
+
+		t.Run("green", func(t *testing.T) {
+			tests := []test{
+				{
+					in:       `<span style="color: red">test</span>`,
+					expected: `<span style="color: red">test</span>`,
+				},
+				{
+					in:       `<span style="color: green">test</span>`,
+					expected: `<span style="color: green">test</span>`,
+				},
+				{
+					in:       `<span style="color: blue">test</span>`,
+					expected: `<span>test</span>`,
+				},
+			}
+
+			for ii, tt := range tests {
+				out := p.Sanitize(tt.in)
+				if out != tt.expected {
+					t.Errorf(
+						"test %d failed;\ninput   : %s\noutput  : %s\nexpected: %s",
+						ii,
+						tt.in,
+						out,
+						tt.expected,
+					)
+				}
+			}
+		})
+
+		p.AllowStyles("color").Matching(regexp.MustCompile("yellow")).OnElements("span")
+
+		t.Run("yellow", func(t *testing.T) {
+			tests := []test{
+				{
+					in:       `<span style="color: red">test</span>`,
+					expected: `<span style="color: red">test</span>`,
+				},
+				{
+					in:       `<span style="color: green">test</span>`,
+					expected: `<span style="color: green">test</span>`,
+				},
+				{
+					in:       `<span style="color: blue">test</span>`,
+					expected: `<span>test</span>`,
+				},
+			}
+
+			for ii, tt := range tests {
+				out := p.Sanitize(tt.in)
+				if out != tt.expected {
+					t.Errorf(
+						"test %d failed;\ninput   : %s\noutput  : %s\nexpected: %s",
+						ii,
+						tt.in,
+						out,
+						tt.expected,
+					)
+				}
+			}
+		})
+	})
+
+	t.Run("AllowURLSchemeWithCustomPolicy", func(t *testing.T) {
+		p := NewPolicy()
+		p.AllowAttrs("href").OnElements("a")
+
+		p.AllowURLSchemeWithCustomPolicy(
+			"http",
+			func(url *url.URL) bool {
+				return url.Hostname() == "example.org"
+			},
+		)
+
+		t.Run("example.org", func(t *testing.T) {
+			tests := []test{
+				{
+					in:       `<a href="http://example.org/">test</a>`,
+					expected: `<a href="http://example.org/">test</a>`,
+				},
+				{
+					in:       `<a href="http://example2.org/">test</a>`,
+					expected: `test`,
+				},
+				{
+					in:       `<a href="http://example4.org/">test</a>`,
+					expected: `test`,
+				},
+			}
+
+			for ii, tt := range tests {
+				out := p.Sanitize(tt.in)
+				if out != tt.expected {
+					t.Errorf(
+						"test %d failed;\ninput   : %s\noutput  : %s\nexpected: %s",
+						ii,
+						tt.in,
+						out,
+						tt.expected,
+					)
+				}
+			}
+		})
+
+		p.AllowURLSchemeWithCustomPolicy(
+			"http",
+			func(url *url.URL) bool {
+				return url.Hostname() == "example2.org"
+			},
+		)
+
+		t.Run("example2.org", func(t *testing.T) {
+			tests := []test{
+				{
+					in:       `<a href="http://example.org/">test</a>`,
+					expected: `<a href="http://example.org/">test</a>`,
+				},
+				{
+					in:       `<a href="http://example2.org/">test</a>`,
+					expected: `<a href="http://example2.org/">test</a>`,
+				},
+				{
+					in:       `<a href="http://example4.org/">test</a>`,
+					expected: `test`,
+				},
+			}
+
+			for ii, tt := range tests {
+				out := p.Sanitize(tt.in)
+				if out != tt.expected {
+					t.Errorf(
+						"test %d failed;\ninput   : %s\noutput  : %s\nexpected: %s",
+						ii,
+						tt.in,
+						out,
+						tt.expected,
+					)
+				}
+			}
+		})
+
+		p.AllowURLSchemeWithCustomPolicy(
+			"http",
+			func(url *url.URL) bool {
+				return url.Hostname() == "example3.org"
+			},
+		)
+
+		t.Run("example3.org", func(t *testing.T) {
+			tests := []test{
+				{
+					in:       `<a href="http://example.org/">test</a>`,
+					expected: `<a href="http://example.org/">test</a>`,
+				},
+				{
+					in:       `<a href="http://example2.org/">test</a>`,
+					expected: `<a href="http://example2.org/">test</a>`,
+				},
+				{
+					in:       `<a href="http://example4.org/">test</a>`,
+					expected: `test`,
+				},
+			}
+
+			for ii, tt := range tests {
+				out := p.Sanitize(tt.in)
+				if out != tt.expected {
+					t.Errorf(
+						"test %d failed;\ninput   : %s\noutput  : %s\nexpected: %s",
+						ii,
+						tt.in,
+						out,
+						tt.expected,
+					)
+				}
+			}
+		})
+	})
+}