Preserve ErrorMessage#getOriginalMessage()

See gh-23417

spring-messaging/src/main/java/org/springframework/messaging/support/MessageBuilder.java

@@ -99,6 +99,7 @@ public final class MessageBuilder<T> {
 		this.headerAccessor.removeHeaders(headerPatterns);
 		return this;
 	}
+
 	/**
 	 * Remove the value for the given header name.
 	 */
@@ -153,7 +154,11 @@ public final class MessageBuilder<T> {
 		}
 		MessageHeaders headersToUse = this.headerAccessor.toMessageHeaders();
 		if (this.payload instanceof Throwable) {
-			return (Message<T>) new ErrorMessage((Throwable) this.payload, headersToUse);
+			Message<?> originalMessage = null;
+			if (this.originalMessage != null && this.originalMessage instanceof ErrorMessage) {
+				originalMessage = ((ErrorMessage) this.originalMessage).getOriginalMessage();
+			}
+			return (Message<T>) new ErrorMessage((Throwable) this.payload, headersToUse, originalMessage);
 		}
 		else {
 			return new GenericMessage<>(this.payload, headersToUse);
@@ -165,6 +170,10 @@ public final class MessageBuilder<T> {
 	 * Create a builder for a new {@link Message} instance pre-populated with all of the
 	 * headers copied from the provided message. The payload of the provided Message will
 	 * also be used as the payload for the new message.
+	 *
+	 * If the provided message is an {@link ErrorMessage} - the
+	 * {@link ErrorMessage#originalMessage} link will be provided to the new instance.
+	 *
 	 * @param message the Message from which the payload and all headers will be copied
 	 */
 	public static <T> MessageBuilder<T> fromMessage(Message<T> message) {

spring-messaging/src/test/java/org/springframework/messaging/support/MessageBuilderTests.java

@@ -16,6 +16,7 @@
 
 package org.springframework.messaging.support;
 
+import java.util.Collections;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.Map;
@@ -107,6 +108,20 @@ public class MessageBuilderTests {
 		assertThat(message2.getHeaders().get("foo")).isEqualTo("bar");
 	}
 
+	@Test
+	public void createErrorMessageFromErrorMessage() {
+		Message<String> originalMessage = MessageBuilder.withPayload("test")
+				.setHeader("foo", "bar").build();
+		RuntimeException errorPayload = new RuntimeException();
+		ErrorMessage errorMessage1 = new ErrorMessage(errorPayload, Collections.singletonMap("baz", "42"), originalMessage);
+		Message<Throwable> errorMessage2 = MessageBuilder.fromMessage(errorMessage1).build();
+		assertThat(errorMessage2).isExactlyInstanceOf(ErrorMessage.class);
+		ErrorMessage actual = (ErrorMessage) errorMessage2;
+		assertThat(actual.getPayload()).isSameAs(errorPayload);
+		assertThat(actual.getHeaders().get("baz")).isEqualTo("42");
+		assertThat(actual.getOriginalMessage()).isSameAs(originalMessage);
+	}
+
 	@Test
 	public void createIdRegenerated() {
 		Message<String> message1 = MessageBuilder.withPayload("test")
@@ -119,20 +134,20 @@ public class MessageBuilderTests {
 	@Test
 	public void testRemove() {
 		Message<Integer> message1 = MessageBuilder.withPayload(1)
-			.setHeader("foo", "bar").build();
+				.setHeader("foo", "bar").build();
 		Message<Integer> message2 = MessageBuilder.fromMessage(message1)
-			.removeHeader("foo")
-			.build();
+				.removeHeader("foo")
+				.build();
 		assertThat(message2.getHeaders().containsKey("foo")).isFalse();
 	}
 
 	@Test
 	public void testSettingToNullRemoves() {
 		Message<Integer> message1 = MessageBuilder.withPayload(1)
-			.setHeader("foo", "bar").build();
+				.setHeader("foo", "bar").build();
 		Message<Integer> message2 = MessageBuilder.fromMessage(message1)
-			.setHeader("foo", null)
-			.build();
+				.setHeader("foo", null)
+				.build();
 		assertThat(message2.getHeaders().containsKey("foo")).isFalse();
 	}
 
@@ -192,7 +207,7 @@ public class MessageBuilderTests {
 
 		assertThatIllegalStateException().isThrownBy(() ->
 				accessor.setHeader("foo", "bar"))
-			.withMessageContaining("Already immutable");
+				.withMessageContaining("Already immutable");
 
 		assertThat(MessageHeaderAccessor.getAccessor(message, MessageHeaderAccessor.class)).isSameAs(accessor);
 	}