Allow schema and validation event handler customization in JAXBContextFactory

jaxb-jakarta/src/main/java/feign/jaxb/JAXBContextFactory.java

@@ -14,6 +14,7 @@
 package feign.jaxb;
 
 import jakarta.xml.bind.*;
+import javax.xml.validation.Schema;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -31,18 +32,44 @@ public final class JAXBContextFactory {
       new ConcurrentHashMap<>(64);
   private final Map<String, Object> properties;
   private final JAXBContextInstantationMode jaxbContextInstantationMode;
+  private final ValidationEventHandler marshallerEventHandler;
+  private final ValidationEventHandler unmarshallerEventHandler;
+  private final Schema marshallerSchema;
+  private final Schema unmashallerSchema;
 
   private JAXBContextFactory(Map<String, Object> properties,
-      JAXBContextInstantationMode jaxbContextInstantationMode) {
+      JAXBContextInstantationMode jaxbContextInstantationMode,
+      ValidationEventHandler marshallerEventHandler,
+      ValidationEventHandler unmarshallerEventHandler,
+      Schema marshallerSchema,
+      Schema unmashallerSchema) {
     this.properties = properties;
     this.jaxbContextInstantationMode = jaxbContextInstantationMode;
+    this.marshallerEventHandler = marshallerEventHandler;
+    this.unmarshallerEventHandler = unmarshallerEventHandler;
+    this.marshallerSchema = marshallerSchema;
+    this.unmashallerSchema = unmashallerSchema;
+  }
+
+  /**
+   * @deprecated please use the constructor with all parameters.
+   */
+  @Deprecated
+  private JAXBContextFactory(Map<String, Object> properties,
+      JAXBContextInstantationMode jaxbContextInstantationMode) {
+    this(properties, jaxbContextInstantationMode, null, null, null, null);
   }
 
   /**
    * Creates a new {@link jakarta.xml.bind.Unmarshaller} that handles the supplied class.
    */
   public Unmarshaller createUnmarshaller(Class<?> clazz) throws JAXBException {
-    return getContext(clazz).createUnmarshaller();
+    Unmarshaller unmarshaller = getContext(clazz).createUnmarshaller();
+    if (unmarshallerEventHandler != null) {
+      unmarshaller.setEventHandler(unmarshallerEventHandler);
+    }
+    unmarshaller.setSchema(unmashallerSchema);
+    return unmarshaller;
   }
 
   /**
@@ -51,6 +78,10 @@ public Unmarshaller createUnmarshaller(Class<?> clazz) throws JAXBException {
   public Marshaller createMarshaller(Class<?> clazz) throws JAXBException {
     Marshaller marshaller = getContext(clazz).createMarshaller();
     setMarshallerProperties(marshaller);
+    if (marshallerEventHandler != null) {
+      marshaller.setEventHandler(marshallerEventHandler);
+    }
+    marshaller.setSchema(marshallerSchema);
     return marshaller;
   }
 
@@ -95,6 +126,14 @@ public static class Builder {
     private JAXBContextInstantationMode jaxbContextInstantationMode =
         JAXBContextInstantationMode.CLASS;
 
+    private ValidationEventHandler marshallerEventHandler;
+
+    private ValidationEventHandler unmarshallerEventHandler;
+
+    private Schema marshallerSchema;
+
+    private Schema unmarshallerSchema;
+
     /**
      * Sets the jaxb.encoding property of any Marshaller created by this factory.
      */
@@ -153,6 +192,38 @@ public Builder withProperty(String key, Object value) {
       return this;
     }
 
+    /**
+     * Sets the validation event handler of any Marshaller created by this factory.
+     */
+    public Builder withMarshallerEventHandler(ValidationEventHandler handler) {
+      this.marshallerEventHandler = handler;
+      return this;
+    }
+
+    /**
+     * Sets the validation event handler of any Unmarshaller created by this factory.
+     */
+    public Builder withUnmarshallerEventHandler(ValidationEventHandler handler) {
+      this.unmarshallerEventHandler = handler;
+      return this;
+    }
+
+    /**
+     * Sets the schema of any Marshaller created by this factory.
+     */
+    public Builder withMarshallerSchema(Schema schema) {
+      this.marshallerSchema = schema;
+      return this;
+    }
+
+    /**
+     * Sets the schema of any Unmarshaller created by this factory.
+     */
+    public Builder withUnmarshallerSchema(Schema schema) {
+      this.unmarshallerSchema = schema;
+      return this;
+    }
+
     /**
      * Provide an instantiation mode for JAXB Contexts, can be class or package, default is class if
      * this method is not called.
@@ -176,7 +247,8 @@ public Builder withJAXBContextInstantiationMode(JAXBContextInstantationMode jaxb
      * Creates a new {@link JAXBContextFactory} instance with a lazy loading cached context
      */
     public JAXBContextFactory build() {
-      return new JAXBContextFactory(properties, jaxbContextInstantationMode);
+      return new JAXBContextFactory(properties, jaxbContextInstantationMode, marshallerEventHandler,
+          unmarshallerEventHandler, marshallerSchema, unmarshallerSchema);
     }
 
     /**
@@ -188,7 +260,7 @@ public JAXBContextFactory build() {
      *         generation most likely due to missing JAXB annotations
      */
     public JAXBContextFactory build(List<Class<?>> classes) throws JAXBException {
-      JAXBContextFactory factory = new JAXBContextFactory(properties, jaxbContextInstantationMode);
+      JAXBContextFactory factory = build();
       factory.preloadContextCache(classes);
       return factory;
     }

jaxb-jakarta/src/test/java/feign/jaxb/JAXBCodecTest.java

@@ -18,18 +18,29 @@
 import feign.RequestTemplate;
 import feign.Response;
 import feign.Util;
+import feign.codec.DecodeException;
+import feign.codec.EncodeException;
 import feign.codec.Encoder;
+import jakarta.xml.bind.MarshalException;
+import jakarta.xml.bind.UnmarshalException;
 import jakarta.xml.bind.annotation.XmlAccessType;
 import jakarta.xml.bind.annotation.XmlAccessorType;
 import jakarta.xml.bind.annotation.XmlElement;
 import jakarta.xml.bind.annotation.XmlRootElement;
+import org.hamcrest.CoreMatchers;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
+import javax.xml.XMLConstants;
+import javax.xml.transform.stream.StreamSource;
+import javax.xml.validation.Schema;
+import javax.xml.validation.SchemaFactory;
+import java.io.StringReader;
 import java.lang.reflect.Type;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Map;
+import java.util.Objects;
 import static feign.Util.UTF_8;
 import static feign.assertj.FeignAssertions.assertThat;
 import static org.junit.Assert.assertEquals;
@@ -256,6 +267,115 @@ public void notFoundDecodesToEmpty() throws Exception {
         .decode(response, byte[].class)).isEmpty();
   }
 
+  @Test
+  public void decodeThrowsExceptionWhenUnmarshallingFailsWithSetSchema() throws Exception {
+    thrown.expect(DecodeException.class);
+    thrown.expectCause(CoreMatchers.instanceOf(UnmarshalException.class));
+    thrown.expectMessage("'Test' is not a valid value for 'integer'.");
+
+    String mockXml = "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"yes\"?><mockIntObject>"
+        + "<value>Test</value></mockIntObject>";
+
+    Response response = Response.builder()
+        .status(200)
+        .reason("OK")
+        .request(Request.create(HttpMethod.GET, "/api", Collections.emptyMap(), null, Util.UTF_8))
+        .headers(Collections.emptyMap())
+        .body(mockXml, UTF_8)
+        .build();
+
+    JAXBContextFactory factory =
+        new JAXBContextFactory.Builder().withUnmarshallerSchema(getMockIntObjSchema()).build();
+    new JAXBDecoder(factory).decode(response, MockIntObject.class);
+  }
+
+  @Test
+  public void decodesIgnoringErrorsWithEventHandler() throws Exception {
+    String mockXml = "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"yes\"?><mockIntObject>"
+        + "<value>Test</value></mockIntObject>";
+
+    Response response = Response.builder()
+        .status(200)
+        .reason("OK")
+        .request(Request.create(HttpMethod.GET, "/api", Collections.emptyMap(), null, Util.UTF_8))
+        .headers(Collections.emptyMap())
+        .body(mockXml, UTF_8)
+        .build();
+
+    JAXBContextFactory factory =
+        new JAXBContextFactory.Builder()
+            .withUnmarshallerSchema(getMockIntObjSchema())
+            .withUnmarshallerEventHandler(event -> true)
+            .build();
+    assertEquals(new MockIntObject(),
+        new JAXBDecoder(factory).decode(response, MockIntObject.class));
+  }
+
+  @Test
+  public void encodeThrowsExceptionWhenMarshallingFailsWithSetSchema() throws Exception {
+    thrown.expect(EncodeException.class);
+    thrown.expectCause(CoreMatchers.instanceOf(MarshalException.class));
+    thrown.expectMessage("The content of element 'mockIntObject' is not complete.");
+
+    JAXBContextFactory jaxbContextFactory = new JAXBContextFactory.Builder()
+        .withMarshallerSchema(getMockIntObjSchema())
+        .build();
+
+    Encoder encoder = new JAXBEncoder(jaxbContextFactory);
+
+    RequestTemplate template = new RequestTemplate();
+    encoder.encode(new MockIntObject(), MockIntObject.class, template);
+  }
+
+  @Test
+  public void encodesIgnoringErrorsWithEventHandler() throws Exception {
+    JAXBContextFactory jaxbContextFactory = new JAXBContextFactory.Builder()
+        .withMarshallerSchema(getMockIntObjSchema())
+        .withMarshallerEventHandler(event -> true)
+        .build();
+
+    Encoder encoder = new JAXBEncoder(jaxbContextFactory);
+
+    RequestTemplate template = new RequestTemplate();
+    encoder.encode(new MockIntObject(), MockIntObject.class, template);
+    assertThat(template).hasBody("<?xml version=\"1.0\" encoding=\"UTF-8\"" +
+        " standalone=\"yes\"?><mockIntObject/>");
+  }
+
+  @XmlRootElement
+  @XmlAccessorType(XmlAccessType.FIELD)
+  static class MockIntObject {
+
+    @XmlElement(required = true)
+    private Integer value;
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o)
+        return true;
+      if (o == null || getClass() != o.getClass())
+        return false;
+      MockIntObject that = (MockIntObject) o;
+      return Objects.equals(value, that.value);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(value);
+    }
+
+  }
+
+  private static Schema getMockIntObjSchema() throws Exception {
+    String schema = "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"yes\"?>\n"
+        + "<xs:schema version=\"1.0\" xmlns:xs=\"http://www.w3.org/2001/XMLSchema\">"
+        + "<xs:element name=\"mockIntObject\" type=\"mockIntObject\"/><xs:complexType name=\"mockIntObject\">"
+        + "<xs:sequence><xs:element name=\"value\" type=\"xs:int\"/></xs:sequence></xs:complexType>"
+        + "</xs:schema>";
+    SchemaFactory schemaFactory = SchemaFactory.newInstance(XMLConstants.W3C_XML_SCHEMA_NS_URI);
+    return schemaFactory.newSchema(new StreamSource(new StringReader(schema)));
+  }
+
   @XmlRootElement
   @XmlAccessorType(XmlAccessType.FIELD)
   static class MockObject {

jaxb-jakarta/src/test/java/feign/jaxb/JAXBContextFactoryTest.java

@@ -16,7 +16,12 @@
 import feign.jaxb.mock.onepackage.AnotherMockedJAXBObject;
 import feign.jaxb.mock.onepackage.MockedJAXBObject;
 import jakarta.xml.bind.Marshaller;
+import jakarta.xml.bind.Unmarshaller;
+import jakarta.xml.bind.ValidationEventHandler;
 import org.junit.Test;
+import javax.xml.XMLConstants;
+import javax.xml.validation.Schema;
+import javax.xml.validation.SchemaFactory;
 import java.lang.reflect.Field;
 import java.util.Arrays;
 import java.util.List;
@@ -75,6 +80,64 @@ public void buildsMarshallerWithFragmentProperty() throws Exception {
     assertTrue((Boolean) marshaller.getProperty(Marshaller.JAXB_FRAGMENT));
   }
 
+  @Test
+  public void buildsMarshallerWithSchema() throws Exception {
+    Schema schema = SchemaFactory.newInstance(XMLConstants.W3C_XML_SCHEMA_NS_URI).newSchema();
+    JAXBContextFactory factory =
+        new JAXBContextFactory.Builder().withMarshallerSchema(schema).build();
+
+    Marshaller marshaller = factory.createMarshaller(Object.class);
+    assertSame(schema, marshaller.getSchema());
+  }
+
+  @Test
+  public void buildsUnmarshallerWithSchema() throws Exception {
+    Schema schema = SchemaFactory.newInstance(XMLConstants.W3C_XML_SCHEMA_NS_URI).newSchema();
+    JAXBContextFactory factory =
+        new JAXBContextFactory.Builder().withUnmarshallerSchema(schema).build();
+
+    Unmarshaller unmarshaller = factory.createUnmarshaller(Object.class);
+    assertSame(schema, unmarshaller.getSchema());
+  }
+
+  @Test
+  public void buildsMarshallerWithCustomEventHandler() throws Exception {
+    ValidationEventHandler handler = event -> false;
+    JAXBContextFactory factory =
+        new JAXBContextFactory.Builder().withMarshallerEventHandler(handler).build();
+
+    Marshaller marshaller = factory.createMarshaller(Object.class);
+    assertSame(handler, marshaller.getEventHandler());
+  }
+
+  @Test
+  public void buildsMarshallerWithDefaultEventHandler() throws Exception {
+    JAXBContextFactory factory =
+        new JAXBContextFactory.Builder().build();
+
+    Marshaller marshaller = factory.createMarshaller(Object.class);
+    assertNotNull(marshaller.getEventHandler());
+  }
+
+  @Test
+  public void buildsUnmarshallerWithCustomEventHandler() throws Exception {
+    ValidationEventHandler handler = event -> false;
+    JAXBContextFactory factory =
+        new JAXBContextFactory.Builder().withUnmarshallerEventHandler(handler).build();
+
+    Unmarshaller unmarshaller = factory.createUnmarshaller(Object.class);
+    assertSame(handler, unmarshaller.getEventHandler());
+  }
+
+  @Test
+  public void buildsUnmarshallerWithDefaultEventHandler() throws Exception {
+    JAXBContextFactory factory =
+        new JAXBContextFactory.Builder().build();
+
+    Unmarshaller unmarshaller = factory.createUnmarshaller(Object.class);
+    assertNotNull(unmarshaller.getEventHandler());
+  }
+
   @Test
   public void testPreloadCache() throws Exception {