GH-1297: Add Jackson PolymorphicTypeValidator to JsonInstanceSerializer

curator-x-discovery/src/main/java/org/apache/curator/x/discovery/details/JsonInstanceSerializer.java

@@ -22,12 +22,20 @@
 import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.JavaType;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.cfg.MapperConfig;
+import com.fasterxml.jackson.databind.jsontype.PolymorphicTypeValidator;
 import com.google.common.annotations.VisibleForTesting;
 import org.apache.curator.x.discovery.ServiceInstance;
 
 /**
  * A serializer that uses Jackson to serialize/deserialize as JSON. IMPORTANT: The instance
  * payload must support Jackson
+ * <p>
+ * The payload is serialized with its class name (see {@link ServiceInstance#getPayload()}). To prevent
+ * deserialization of arbitrary classes, by default only the payload class and its subtypes are
+ * accepted when deserializing the payload. Polymorphic types nested within the payload whose declared
+ * base type is not {@link Object} are not restricted. Use {@link #JsonInstanceSerializer(Class, boolean, PolymorphicTypeValidator)}
+ * to supply a custom {@link PolymorphicTypeValidator}.
  */
 public class JsonInstanceSerializer<T> implements InstanceSerializer<T> {
     private final ObjectMapper mapper;
@@ -66,16 +74,69 @@ public JsonInstanceSerializer(Class<T> payloadClass, boolean compatibleSerializa
         this(payloadClass, compatibleSerializationMode, false);
     }
 
+    /**
+     * Same as {@link #JsonInstanceSerializer(Class, boolean)} but with a custom {@link PolymorphicTypeValidator}
+     * that limits which classes may be deserialized as the payload.
+     *
+     * @param payloadClass used to validate payloads when deserializing
+     * @param compatibleSerializationMode pass true to serialize in a manner that supports clients pre-CURATOR-275
+     * @param polymorphicTypeValidator validates the payload class named in the serialized JSON
+     */
+    public JsonInstanceSerializer(
+            Class<T> payloadClass,
+            boolean compatibleSerializationMode,
+            PolymorphicTypeValidator polymorphicTypeValidator) {
+        this(payloadClass, compatibleSerializationMode, false, polymorphicTypeValidator);
+    }
+
     @VisibleForTesting
     JsonInstanceSerializer(
             Class<T> payloadClass, boolean compatibleSerializationMode, boolean failOnUnknownProperties) {
+        this(payloadClass, compatibleSerializationMode, failOnUnknownProperties, defaultValidator(payloadClass));
+    }
+
+    private JsonInstanceSerializer(
+            Class<T> payloadClass,
+            boolean compatibleSerializationMode,
+            boolean failOnUnknownProperties,
+            PolymorphicTypeValidator polymorphicTypeValidator) {
         this.payloadClass = payloadClass;
         this.compatibleSerializationMode = compatibleSerializationMode;
         mapper = new ObjectMapper();
+        mapper.setPolymorphicTypeValidator(polymorphicTypeValidator);
         mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, failOnUnknownProperties);
         type = mapper.getTypeFactory().constructType(ServiceInstance.class);
     }
 
+    private static PolymorphicTypeValidator defaultValidator(Class<?> payloadClass) {
+        return new PayloadTypeValidator(payloadClass);
+    }
+
+    /**
+     * The payload is declared as {@link Object} (due to type erasure) so this validator only restricts
+     * polymorphic values with an {@link Object} base type, which must be the payload class or a subtype of it.
+     * Polymorphic values with any other base type are allowed (as they were before this validator was added).
+     */
+    private static class PayloadTypeValidator extends PolymorphicTypeValidator.Base {
+        private static final long serialVersionUID = 1L;
+
+        private final Class<?> payloadClass;
+
+        PayloadTypeValidator(Class<?> payloadClass) {
+            this.payloadClass = payloadClass;
+        }
+
+        @Override
+        public Validity validateBaseType(MapperConfig<?> config, JavaType baseType) {
+            return baseType.hasRawClass(Object.class) ? Validity.INDETERMINATE : Validity.ALLOWED;
+        }
+
+        @Override
+        public Validity validateSubType(MapperConfig<?> config, JavaType baseType, JavaType subType) {
+            return payloadClass.isAssignableFrom(subType.getRawClass()) ? Validity.ALLOWED : Validity.DENIED;
+        }
+    }
+
     @SuppressWarnings({"unchecked"})
     @Override
     public ServiceInstance<T> deserialize(byte[] bytes) throws Exception {

curator-x-discovery/src/test/java/org/apache/curator/x/discovery/TestJsonInstanceSerializer.java

@@ -20,8 +20,12 @@
 package org.apache.curator.x.discovery;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
+import com.fasterxml.jackson.databind.exc.InvalidTypeIdException;
+import com.fasterxml.jackson.databind.jsontype.BasicPolymorphicTypeValidator;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -157,6 +161,107 @@ public void testPayloadClass() throws Exception {
         assertEquals(instance.isEnabled(), rhs.isEnabled());
     }
 
+    @Test
+    public void testDisallowedPayloadClass() throws Exception {
+        JsonInstanceSerializer<Object> objectSerializer = new JsonInstanceSerializer<Object>(Object.class);
+        byte[] bytes = objectSerializer.serialize(new ServiceInstance<Object>(
+                "name",
+                "id",
+                "address",
+                10,
+                20,
+                new OtherPayload("other"),
+                0,
+                ServiceType.DYNAMIC,
+                new UriSpec("{a}/b/{c}"),
+                true));
+
+        JsonInstanceSerializer<Payload> serializer = new JsonInstanceSerializer<Payload>(Payload.class);
+        assertThrows(InvalidTypeIdException.class, () -> serializer.deserialize(bytes));
+    }
+
+    @Test
+    public void testCustomPolymorphicTypeValidator() throws Exception {
+        JsonInstanceSerializer<Object> serializer = new JsonInstanceSerializer<Object>(
+                Object.class,
+                false,
+                BasicPolymorphicTypeValidator.builder()
+                        .allowIfSubType(Payload.class)
+                        .build());
+
+        ServiceInstance<Object> allowed = new ServiceInstance<Object>(
+                "name",
+                "id",
+                "address",
+                10,
+                20,
+                new Payload("test"),
+                0,
+                ServiceType.DYNAMIC,
+                new UriSpec("{a}/b/{c}"),
+                true);
+        assertEquals(
+                allowed.getPayload(),
+                serializer.deserialize(serializer.serialize(allowed)).getPayload());
+
+        byte[] disallowedBytes = serializer.serialize(new ServiceInstance<Object>(
+                "name",
+                "id",
+                "address",
+                10,
+                20,
+                new OtherPayload("other"),
+                0,
+                ServiceType.DYNAMIC,
+                new UriSpec("{a}/b/{c}"),
+                true));
+        assertThrows(InvalidTypeIdException.class, () -> serializer.deserialize(disallowedBytes));
+    }
+
+    @Test
+    public void testPayloadWithNestedPolymorphicField() throws Exception {
+        JsonInstanceSerializer<NestedPayload> serializer =
+                new JsonInstanceSerializer<NestedPayload>(NestedPayload.class);
+        NestedPayload payload = new NestedPayload();
+        payload.setNested(new Payload("nested"));
+        ServiceInstance<NestedPayload> instance = new ServiceInstance<NestedPayload>(
+                "name", "id", "address", 10, 20, payload, 0, ServiceType.DYNAMIC, new UriSpec("{a}/b/{c}"), true);
+
+        ServiceInstance<NestedPayload> rhs = serializer.deserialize(serializer.serialize(instance));
+        assertEquals(payload.getNested(), rhs.getPayload().getNested());
+    }
+
+    public static class NestedPayload {
+        private Payload nested;
+
+        @JsonTypeInfo(use = JsonTypeInfo.Id.CLASS)
+        public Payload getNested() {
+            return nested;
+        }
+
+        public void setNested(Payload nested) {
+            this.nested = nested;
+        }
+    }
+
+    public static class OtherPayload {
+        private String val;
+
+        public OtherPayload() {}
+
+        public OtherPayload(String val) {
+            this.val = val;
+        }
+
+        public String getVal() {
+            return val;
+        }
+
+        public void setVal(String val) {
+            this.val = val;
+        }
+    }
+
     public static class Payload {
         private String val;