reworked code that munges maps when compiling element and attribute

policies to have fewer special cases and allow styling policies to compose.

Author: mikesamuel

src/main/java/org/owasp/html/HtmlPolicyBuilder.java

@@ -638,22 +638,36 @@ public PolicyFactory toFactory() {
  textContainerSet.add(textContainer.getKey());
  }
  }
- return new PolicyFactory(compilePolicies(), textContainerSet.build(),
- ImmutableMap.copyOf(globalAttrPolicies),
- preprocessor, postprocessor);
+ CompiledState compiled = compilePolicies();
+
+ return new PolicyFactory(
+ compiled.compiledPolicies, textContainerSet.build(),
+ ImmutableMap.copyOf(compiled.globalAttrPolicies),
+ preprocessor, postprocessor);
  }
 
  // Speed up subsequent builds by caching the compiled policies.
- private transient ImmutableMap<String, ElementAndAttributePolicies>
- compiledPolicies;
+ private transient CompiledState compiledState;
+
+ private static final class CompiledState {
+ final Map<String, AttributePolicy> globalAttrPolicies;
+ final ImmutableMap<String, ElementAndAttributePolicies> compiledPolicies;
+
+ CompiledState(
+ Map<String, AttributePolicy> globalAttrPolicies,
+ ImmutableMap<String, ElementAndAttributePolicies> compiledPolicies) {
+ this.globalAttrPolicies = globalAttrPolicies;
+ this.compiledPolicies = compiledPolicies;
+ }
+ }
 
  /** Called by mutators to signal that any compiled policy is out-of-date. */
  private void invalidateCompiledState() {
- compiledPolicies = null;
+ compiledState = null;
  }
 
- private ImmutableMap<String, ElementAndAttributePolicies> compilePolicies() {
- if (compiledPolicies != null) { return compiledPolicies; }
+ private CompiledState compilePolicies() {
+ if (compiledState != null) { return compiledState; }
 
  // Copy maps before normalizing in case builder is reused.
  @SuppressWarnings("hiding")
@@ -741,6 +755,13 @@ public String apply(String url) {
  }
  }
  }
+ if (stylingPolicy != null) {
+ AttributePolicy old = globalAttrPolicies.get("style");
+ if (old != null) {
+ globalAttrPolicies.put(
+ "style", AttributePolicy.Util.join(old, stylingPolicy));
+ }
+ }
 
  ImmutableMap.Builder<String, ElementAndAttributePolicies> policiesBuilder
  = ImmutableMap.builder();
@@ -753,18 +774,29 @@ public String apply(String url) {
 
  Map<String, AttributePolicy> elAttrPolicies
  = attrPolicies.get(elementName);
- if (elAttrPolicies == null) { elAttrPolicies = ImmutableMap.of(); }
+ if (elAttrPolicies == null) {
+ elAttrPolicies = ImmutableMap.of();
+ }
+ if (stylingPolicy != null) {
+ AttributePolicy old = elAttrPolicies.get("style");
+ if (old != null) {
+ attrPolicies.put(
+ elementName,
+ elAttrPolicies = Maps.newLinkedHashMap(elAttrPolicies));
+ elAttrPolicies.put(
+ "style", AttributePolicy.Util.join(old, stylingPolicy));
+ }
+ }
+
  ImmutableMap.Builder<String, AttributePolicy> attrs
  = ImmutableMap.builder();
+
  for (Map.Entry<String, AttributePolicy> ape : elAttrPolicies.entrySet()) {
  String attributeName = ape.getKey();
  // Handle below so we don't end up putting the same key into the map
  // twice. ImmutableMap.Builder hates that.
  if (globalAttrPolicies.containsKey(attributeName)) { continue; }
  AttributePolicy policy = ape.getValue();
- if ("style".equals(attributeName)) {
- policy = AttributePolicy.Util.join(policy, stylingPolicy);
- }
  if (!AttributePolicy.REJECT_ALL_ATTRIBUTE_POLICY.equals(policy)) {
  attrs.put(attributeName, policy);
  }
@@ -774,9 +806,6 @@ public String apply(String url) {
  String attributeName = ape.getKey();
  AttributePolicy policy = AttributePolicy.Util.join(
  elAttrPolicies.get(attributeName), ape.getValue());
- if ("style".equals(attributeName)) {
- policy = AttributePolicy.Util.join(policy, stylingPolicy);
- }
  if (!AttributePolicy.REJECT_ALL_ATTRIBUTE_POLICY.equals(policy)) {
  attrs.put(attributeName, policy);
  }
@@ -788,7 +817,7 @@ public String apply(String url) {
  elementName,
  elPolicy, attrs.build(), skipIfEmpty.contains(elementName)));
  }
- return compiledPolicies = policiesBuilder.build();
+ return new CompiledState(globalAttrPolicies, policiesBuilder.build());
  }
 
  /**

src/main/java/org/owasp/html/StylingPolicy.java

@@ -32,8 +32,11 @@
 
 import javax.annotation.Nullable;
 
+import org.owasp.html.JoinedAttributePolicy.JoinableAttributePolicy;
+
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
+import com.google.common.base.Functions;
 import com.google.common.collect.Lists;
 
 /**
@@ -42,7 +45,7 @@
  * ones to reduce the attack-surface.
  */
 @TCB
-final class StylingPolicy implements AttributePolicy {
+final class StylingPolicy implements JoinableAttributePolicy {
 
  final CssSchema cssSchema;
  final Function<String, String> urlRewriter;
@@ -258,4 +261,32 @@ public boolean equals(Object o) {
  public int hashCode() {
  return cssSchema.hashCode();
  }
+
+ public Joinable.JoinStrategy<JoinableAttributePolicy> getJoinStrategy() {
+ return StylingPolicyJoinStrategy.INSTANCE;
+ }
+
+ static final class StylingPolicyJoinStrategy
+ implements Joinable.JoinStrategy<JoinableAttributePolicy> {
+ static final StylingPolicyJoinStrategy INSTANCE =
+ new StylingPolicyJoinStrategy();
+
+ public JoinableAttributePolicy join(
+ Iterable<? extends JoinableAttributePolicy> toJoin) {
+ Function<String, String> identity = Functions.<String>identity();
+ CssSchema cssSchema = null;
+ Function<String, String> urlRewriter = identity;
+ for (JoinableAttributePolicy p : toJoin) {
+ StylingPolicy sp = (StylingPolicy) p;
+ cssSchema = cssSchema == null
+ ? sp.cssSchema : CssSchema.union(cssSchema, sp.cssSchema);
+ urlRewriter = urlRewriter.equals(identity)
+ || urlRewriter.equals(sp.urlRewriter)
+ ? sp.urlRewriter
+ : Functions.compose(urlRewriter, sp.urlRewriter);
+ }
+ return new StylingPolicy(cssSchema, urlRewriter);
+ }
+
+ }
 }

src/test/java/org/owasp/html/SanitizersTest.java

@@ -319,12 +319,12 @@ public static final void testScriptInTable() {
  @Test
  public static final void testAndOrdering() {
  String input = ""
- + "xss<a href=\"http://www.google.de\" style=\"color:red;\""
+ + "xss<a href=\"http://www.google.de\" style=\"color:red\""
  + " onmouseover=alert(1) onmousemove=\"alert(2)\" onclick=alert(3)>"
  + "g"
  + "<img src=\"http://example.org\"/>oogle</a>";
  String want = ""
- + "xss<a href=\"http://www.google.de\" style=\"color:red;\""
+ + "xss<a href=\"http://www.google.de\" style=\"color:red\""
  + " rel=\"nofollow\">"
  + "g"
  + "<img src=\"http://example.org\" />oogle</a>";