Refactoring ConstraintValidators to sanitize violation message templates

titus-api/src/main/java/com/netflix/titus/api/jobmanager/model/job/sanitizer/SchedulingConstraintSetValidator.java

@@ -22,18 +22,20 @@
 import java.lang.annotation.Target;
 import java.util.HashSet;
 import java.util.Set;
+import java.util.function.Function;
 import javax.validation.Constraint;
-import javax.validation.ConstraintValidator;
 import javax.validation.ConstraintValidatorContext;
 import javax.validation.Payload;
 
 import com.netflix.titus.api.jobmanager.model.job.Container;
+import com.netflix.titus.common.model.sanitizer.internal.AbstractConstraintValidator;
 
 import static java.lang.annotation.ElementType.TYPE;
 
 /**
+ *
  */
-public class SchedulingConstraintSetValidator implements ConstraintValidator<SchedulingConstraintSetValidator.SchedulingConstraintSet, Container> {
+public class SchedulingConstraintSetValidator extends AbstractConstraintValidator<SchedulingConstraintSetValidator.SchedulingConstraintSet, Container> {
 
     @Target({TYPE})
     @Retention(RetentionPolicy.RUNTIME)
@@ -53,7 +55,7 @@ public void initialize(SchedulingConstraintSet constraintAnnotation) {
     }
 
     @Override
-    public boolean isValid(Container container, ConstraintValidatorContext context) {
+    protected boolean isValid(Container container, Function<String, ConstraintValidatorContext.ConstraintViolationBuilder> constraintViolationBuilderFunction) {
         if (container == null) {
             return true;
         }
@@ -64,7 +66,7 @@ public boolean isValid(Container container, ConstraintValidatorContext context)
             return true;
         }
 
-        context.buildConstraintViolationWithTemplate(
+        constraintViolationBuilderFunction.apply(
                 "Soft and hard constraints not unique. Shared constraints: " + common
         ).addConstraintViolation().disableDefaultConstraintViolation();
         return false;

titus-api/src/main/java/com/netflix/titus/api/jobmanager/model/job/sanitizer/SchedulingConstraintValidator.java

@@ -22,22 +22,22 @@
 import java.lang.annotation.Target;
 import java.util.HashSet;
 import java.util.Map;
-import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import javax.validation.Constraint;
-import javax.validation.ConstraintValidator;
 import javax.validation.ConstraintValidatorContext;
 import javax.validation.Payload;
 
 import com.netflix.titus.api.jobmanager.JobConstraints;
+import com.netflix.titus.common.model.sanitizer.internal.AbstractConstraintValidator;
 
 import static java.lang.annotation.ElementType.ANNOTATION_TYPE;
 import static java.lang.annotation.ElementType.CONSTRUCTOR;
 import static java.lang.annotation.ElementType.FIELD;
 import static java.lang.annotation.ElementType.METHOD;
 import static java.lang.annotation.ElementType.PARAMETER;
 
-public class SchedulingConstraintValidator implements ConstraintValidator<SchedulingConstraintValidator.SchedulingConstraint, Map<String, String>> {
+public class SchedulingConstraintValidator extends AbstractConstraintValidator<SchedulingConstraintValidator.SchedulingConstraint, Map<String, String>> {
 
     @Target({METHOD, FIELD, ANNOTATION_TYPE, CONSTRUCTOR, PARAMETER})
     @Retention(RetentionPolicy.RUNTIME)
@@ -69,14 +69,13 @@ public void initialize(SchedulingConstraint constraintAnnotation) {
     }
 
     @Override
-    public boolean isValid(Map<String, String> value, ConstraintValidatorContext context) {
-        Set<String> namesInLowerCase = value.keySet().stream().map(String::toLowerCase).collect(Collectors.toSet());
-        HashSet<String> unknown = new HashSet<>(namesInLowerCase);
+    protected boolean isValid(Map<String, String> value, Function<String, ConstraintValidatorContext.ConstraintViolationBuilder> constraintViolationBuilderFunction) {
+        HashSet<String> unknown = value.keySet().stream().map(String::toLowerCase).collect(Collectors.toCollection(HashSet::new));
         unknown.removeAll(JobConstraints.CONSTRAINT_NAMES);
         if (unknown.isEmpty()) {
             return true;
         }
-        context.buildConstraintViolationWithTemplate("Unrecognized constraints " + unknown)
+        constraintViolationBuilderFunction.apply("Unrecognized constraints " + unknown)
                 .addConstraintViolation().disableDefaultConstraintViolation();
         return false;
     }

titus-api/src/test/java/com/netflix/titus/api/jobmanager/model/job/sanitizer/SchedulingConstraintValidatorTest.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2020 Netflix, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.netflix.titus.api.jobmanager.model.job.sanitizer;
+
+import com.google.common.collect.ImmutableMap;
+import com.netflix.titus.api.jobmanager.JobConstraints;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import javax.validation.ConstraintViolation;
+import javax.validation.Validation;
+import javax.validation.Validator;
+import java.util.*;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+
+public class SchedulingConstraintValidatorTest {
+
+    private Validator validator;
+
+    private static final String EL_EXPRESSION_MESSAGE = "#{#this.class.name.substring(0,5) == 'com.g' ? 'FOO' : T(java.lang.Runtime).getRuntime().exec(new java.lang.String(T(java.util.Base64).getDecoder().decode('dG91Y2ggL3RtcC9wd25lZA=='))).class.name}";
+    private static final String RCE_VIOLATION_MESSAGE = "Unrecognized constraints [\\#\\{\\#this.class.name.substring(0,5) == 'com.g' ? 'foo' : t(java.lang.runtime).getruntime().exec(new java.lang.string(t(java.util.base64).getdecoder().decode('dg91y2ggl3rtcc9wd25lza=='))).class.name\\}]";
+
+    @Before
+    public void setUp() {
+        validator = Validation.buildDefaultValidatorFactory().getValidator();
+    }
+
+    @Test
+    public void testValidSchedulingConstraint() {
+        Map<String, String> validConstraintMap = JobConstraints.CONSTRAINT_NAMES.stream().collect(Collectors.toMap(Function.identity(), containerName -> "random_string"));
+        assertThat(validator.validate(new ConstraintMap(validConstraintMap))).isEmpty();
+    }
+
+    @Test
+    public void testInvalidSchedulingConstraint() {
+        String randomKey = UUID.randomUUID().toString();
+        String expectedViolationString = "Unrecognized constraints [" + randomKey + "]";
+        Map<String, String> invalidConstraintMap = ImmutableMap.of(randomKey, "test");
+        Set<ConstraintViolation<ConstraintMap>> violations = validator.validate(new ConstraintMap(invalidConstraintMap));
+        ConstraintViolation<ConstraintMap> constraintMapConstraintViolation = violations.stream().findFirst().orElseThrow(() -> new AssertionError("Expected violation not found"));
+        assertThat(constraintMapConstraintViolation.getMessage()).isEqualTo(expectedViolationString);
+    }
+
+    @Test
+    public void testRceAttemptInSchedulingConstraint() {
+        Map<String, String> invalidConstraintMap = ImmutableMap.of(EL_EXPRESSION_MESSAGE, "test");
+        Set<ConstraintViolation<ConstraintMap>> violations = validator.validate(new ConstraintMap(invalidConstraintMap));
+        ConstraintViolation<ConstraintMap> constraintMapConstraintViolation = violations.stream().findFirst().orElseThrow(() -> new AssertionError("Expected violation not found"));
+        assertThat(constraintMapConstraintViolation.getMessageTemplate()).isEqualTo(RCE_VIOLATION_MESSAGE);
+    }
+
+    static class ConstraintMap {
+
+        @SchedulingConstraintValidator.SchedulingConstraint
+        final Map<String, String> constraintMap;
+
+        public ConstraintMap(Map<String, String> map) {
+            this.constraintMap = map;
+        }
+    }
+}

titus-common/src/main/java/com/netflix/titus/common/model/sanitizer/internal/AbstractConstraintValidator.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2020 Netflix, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.netflix.titus.common.model.sanitizer.internal;
+
+import javax.validation.ConstraintValidator;
+import javax.validation.ConstraintValidatorContext;
+import java.lang.annotation.Annotation;
+import java.util.function.Function;
+
+public abstract class AbstractConstraintValidator<A extends Annotation, T> implements ConstraintValidator<A, T> {
+
+    /**
+     * Escape all special characters that participate in EL expressions so the the message string
+     * cannot be classified as a template for interpolation.
+     *
+     * @param message string that needs to be sanitized
+     * @return copy of the input string with '{','}','#' and '$' characters escaped
+     */
+    private static String sanitizeMessage(String message) {
+        return message.replaceAll("([}{$#])", "\\\\$1");
+    }
+
+    @Override
+    final public boolean isValid(T type, ConstraintValidatorContext context) {
+        return this.isValid(type, message -> {
+            String sanitizedMessage = sanitizeMessage(message);
+            return context.buildConstraintViolationWithTemplate(sanitizedMessage);
+        });
+    }
+
+    /**
+     * Implementing classes will need to apply the builder function with the violation message string
+     * to retrieve the underlying instance of {@link javax.validation.ConstraintValidatorContext} in order
+     * to continue add any violations.
+     *
+     * @param type                               type of the object under validation
+     * @param constraintViolationBuilderFunction function to apply with a violation message string
+     * @return validation status
+     */
+    abstract protected boolean isValid(T type, Function<String, ConstraintValidatorContext.ConstraintViolationBuilder> constraintViolationBuilderFunction);
+
+}

titus-common/src/main/java/com/netflix/titus/common/model/sanitizer/internal/CollectionValidator.java

@@ -21,13 +21,14 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.TreeSet;
+import java.util.function.Function;
 import java.util.stream.Collectors;
-import javax.validation.ConstraintValidator;
+
 import javax.validation.ConstraintValidatorContext;
 
 import com.netflix.titus.common.model.sanitizer.CollectionInvariants;
 
-public class CollectionValidator implements ConstraintValidator<CollectionInvariants, Object> {
+public class CollectionValidator extends AbstractConstraintValidator<CollectionInvariants, Object> {
 
     private CollectionInvariants constraintAnnotation;
 
@@ -37,49 +38,50 @@ public void initialize(CollectionInvariants constraintAnnotation) {
     }
 
     @Override
-    public boolean isValid(Object value, ConstraintValidatorContext context) {
+    protected boolean isValid(Object value, Function<String, ConstraintValidatorContext.ConstraintViolationBuilder> constraintViolationBuilderFunction) {
+
         if (value == null) {
             return true;
         }
         if (value instanceof Collection) {
-            return isValid((Collection<?>) value, context);
+            return isValid((Collection<?>) value, constraintViolationBuilderFunction);
         }
         if (value instanceof Map) {
-            return isValid((Map<?, ?>) value, context);
+            return isValid((Map<?, ?>) value, constraintViolationBuilderFunction);
         }
         return false;
     }
 
-    private boolean isValid(Collection<?> value, ConstraintValidatorContext context) {
+    private boolean isValid(Collection<?> value, Function<String, ConstraintValidatorContext.ConstraintViolationBuilder> constraintViolationBuilderFunction) {
         if (value.isEmpty()) {
             return true;
         }
 
         if (!constraintAnnotation.allowNullValues()) {
             if (value.stream().anyMatch(Objects::isNull)) {
-                attachMessage(context, "null values not allowed");
+                attachMessage(constraintViolationBuilderFunction, "null values not allowed");
                 return false;
             }
         }
 
         return true;
     }
 
-    private boolean isValid(Map<?, ?> value, ConstraintValidatorContext context) {
+    private boolean isValid(Map<?, ?> value, Function<String, ConstraintValidatorContext.ConstraintViolationBuilder> constraintViolationBuilderFunction) {
         if (value.isEmpty()) {
             return true;
         }
 
         if (!constraintAnnotation.allowEmptyKeys()) {
             if (value.keySet().stream().anyMatch(key -> key == null || (key instanceof String && ((String) key).isEmpty()))) {
-                attachMessage(context, "empty key names not allowed");
+                attachMessage(constraintViolationBuilderFunction, "empty key names not allowed");
                 return false;
             }
         }
 
         if (!constraintAnnotation.allowNullKeys()) {
             if (value.keySet().stream().anyMatch(Objects::isNull)) {
-                attachMessage(context, "null key names not allowed");
+                attachMessage(constraintViolationBuilderFunction, "null key names not allowed");
                 return false;
             }
         }
@@ -90,16 +92,15 @@ private boolean isValid(Map<?, ?> value, ConstraintValidatorContext context) {
                     .map(e -> e.getKey() instanceof String ? (String) e.getKey() : "<not_string>")
                     .collect(Collectors.toSet());
             if (!badEntryKeys.isEmpty()) {
-                attachMessage(context, "null values found for keys: " + new TreeSet<>(badEntryKeys));
+                attachMessage(constraintViolationBuilderFunction, "null values found for keys: " + new TreeSet<>(badEntryKeys));
                 return false;
             }
         }
 
         return true;
     }
 
-    private void attachMessage(ConstraintValidatorContext context, String message) {
-        context.buildConstraintViolationWithTemplate(message).addConstraintViolation();
-        context.disableDefaultConstraintViolation();
+    private void attachMessage(Function<String, ConstraintValidatorContext.ConstraintViolationBuilder> constraintViolationBuilderFunction, String message) {
+        constraintViolationBuilderFunction.apply(message).addConstraintViolation().disableDefaultConstraintViolation();
     }
 }

titus-common/src/main/java/com/netflix/titus/common/model/sanitizer/internal/NeverNullValidator.java

@@ -20,26 +20,26 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
-import javax.validation.ConstraintValidator;
+import java.util.function.Function;
+
 import javax.validation.ConstraintValidatorContext;
 
 import com.netflix.titus.common.model.sanitizer.ClassFieldsNotNull;
 import com.netflix.titus.common.util.ReflectionExt;
 
-public class NeverNullValidator implements ConstraintValidator<ClassFieldsNotNull, Object> {
+public class NeverNullValidator extends AbstractConstraintValidator<ClassFieldsNotNull, Object> {
 
     @Override
     public void initialize(ClassFieldsNotNull constraintAnnotation) {
     }
 
     @Override
-    public boolean isValid(Object value, ConstraintValidatorContext context) {
+    protected boolean isValid(Object value, Function<String, ConstraintValidatorContext.ConstraintViolationBuilder> constraintViolationBuilderFunction) {
         Set<String> nullFields = validate(value);
         if (nullFields.isEmpty()) {
             return true;
         }
-        context.buildConstraintViolationWithTemplate(buildMessage(nullFields)).addConstraintViolation();
-        context.disableDefaultConstraintViolation();
+        constraintViolationBuilderFunction.apply(buildMessage(nullFields)).addConstraintViolation().disableDefaultConstraintViolation();
         return false;
     }