From d7970e4ab82484fea178d87f7664056b27c704e1 Mon Sep 17 00:00:00 2001
From: Arjen Poutsma <poutsmaa@vmware.com>
Date: Thu, 8 Jun 2023 15:29:56 +0200
Subject: [PATCH] Support JAXBElement in Jaxb2XmlEncoder

This commit introduces support for JAXBElements in the Jaxb2XmlEncoder.

Closes gh-30552
---
 .../http/codec/xml/Jaxb2XmlEncoder.java           | 14 ++++++++++++--
 .../http/codec/xml/Jaxb2XmlEncoderTests.java      | 15 +++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/spring-web/src/main/java/org/springframework/http/codec/xml/Jaxb2XmlEncoder.java b/spring-web/src/main/java/org/springframework/http/codec/xml/Jaxb2XmlEncoder.java
index 678e6c64ae70..482909528c71 100644
--- a/spring-web/src/main/java/org/springframework/http/codec/xml/Jaxb2XmlEncoder.java
+++ b/spring-web/src/main/java/org/springframework/http/codec/xml/Jaxb2XmlEncoder.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@
 import java.util.Map;
 import java.util.function.Function;
 
+import jakarta.xml.bind.JAXBElement;
 import jakarta.xml.bind.JAXBException;
 import jakarta.xml.bind.MarshalException;
 import jakarta.xml.bind.Marshaller;
@@ -121,7 +122,7 @@ public DataBuffer encodeValue(Object value, DataBufferFactory bufferFactory,
 		DataBuffer buffer = bufferFactory.allocateBuffer(1024);
 		try {
 			OutputStream outputStream = buffer.asOutputStream();
-			Class<?> clazz = ClassUtils.getUserClass(value);
+			Class<?> clazz = getMarshallerType(value);
 			Marshaller marshaller = initMarshaller(clazz);
 			marshaller.marshal(value, outputStream);
 			release = false;
@@ -140,6 +141,15 @@ public DataBuffer encodeValue(Object value, DataBufferFactory bufferFactory,
 		}
 	}
 
+	private static Class<?> getMarshallerType(Object value) {
+		if (value instanceof JAXBElement<?> jaxbElement) {
+			return jaxbElement.getDeclaredType();
+		}
+		else {
+			return ClassUtils.getUserClass(value);
+		}
+	}
+
 	private Marshaller initMarshaller(Class<?> clazz) throws CodecException, JAXBException {
 		Marshaller marshaller = this.jaxbContexts.createMarshaller(clazz);
 		marshaller.setProperty(Marshaller.JAXB_ENCODING, StandardCharsets.UTF_8.name());
diff --git a/spring-web/src/test/java/org/springframework/http/codec/xml/Jaxb2XmlEncoderTests.java b/spring-web/src/test/java/org/springframework/http/codec/xml/Jaxb2XmlEncoderTests.java
index ad38d40208f2..e134ed57cea3 100644
--- a/spring-web/src/test/java/org/springframework/http/codec/xml/Jaxb2XmlEncoderTests.java
+++ b/spring-web/src/test/java/org/springframework/http/codec/xml/Jaxb2XmlEncoderTests.java
@@ -20,6 +20,9 @@
 import java.util.List;
 import java.util.function.Consumer;
 
+import javax.xml.namespace.QName;
+
+import jakarta.xml.bind.JAXBElement;
 import jakarta.xml.bind.annotation.XmlElement;
 import jakarta.xml.bind.annotation.XmlElements;
 import jakarta.xml.bind.annotation.XmlRootElement;
@@ -76,6 +79,18 @@ public void encode() {
 				.verifyComplete());
 	}
 
+	@Test
+	public void encodeJaxbElement() {
+		Mono<JAXBElement<Pojo>> input = Mono.just(new JAXBElement<>(new QName("baz"), Pojo.class,
+				new Pojo("foofoo", "barbar")));
+
+		testEncode(input, Pojo.class, step -> step
+				.consumeNextWith(expectXml(
+						"<?xml version='1.0' encoding='UTF-8' standalone='yes'?>" +
+								"<baz><bar>barbar</bar><foo>foofoo</foo></baz>"))
+				.verifyComplete());
+	}
+
 	@Test
 	public void encodeError() {
 		Flux<Pojo> input = Flux.error(RuntimeException::new);