From 2c47b2b6fae664e31a21c0d07fa255a3a5890f1d Mon Sep 17 00:00:00 2001
From: Mitchell <38743745+mitchmcd18@users.noreply.github.com>
Date: Fri, 9 Aug 2024 05:16:49 +1200
Subject: [PATCH] GH-9369: Fix MutableMessageHeaders for serialization

Fixes: #9369

* Implement `MutableMessageHeaders.readResolve()` to reinstate the instance

(cherry picked from commit 02b58277f49cc9561992c449431033870e08156c)
---
 .../support/MutableMessageHeaders.java        |  8 ++++
 .../support/MutableMessageTests.java          | 37 +++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/spring-integration-core/src/main/java/org/springframework/integration/support/MutableMessageHeaders.java b/spring-integration-core/src/main/java/org/springframework/integration/support/MutableMessageHeaders.java
index 1730050712..eda3b895da 100644
--- a/spring-integration-core/src/main/java/org/springframework/integration/support/MutableMessageHeaders.java
+++ b/spring-integration-core/src/main/java/org/springframework/integration/support/MutableMessageHeaders.java
@@ -16,6 +16,8 @@
 
 package org.springframework.integration.support;
 
+import java.io.ObjectStreamException;
+import java.io.Serial;
 import java.nio.ByteBuffer;
 import java.util.Map;
 import java.util.UUID;
@@ -32,6 +34,7 @@ import org.springframework.messaging.MessageHeaders;
  * @author David Turanski
  * @author Artem Bilan
  * @author Nathan Kurtyka
+ * @author Mitchell McDonald
  *
  * @since 4.2
  */
@@ -74,6 +77,11 @@ public class MutableMessageHeaders extends MessageHeaders {
 		return super.getRawHeaders().remove(key);
 	}
 
+	@Serial
+	private Object readResolve() throws ObjectStreamException {
+		return new MutableMessageHeaders(this);
+	}
+
 	@Nullable
 	private static UUID extractId(@Nullable Map<String, Object> headers) {
 		if (headers != null && headers.containsKey(MessageHeaders.ID)) {
diff --git a/spring-integration-core/src/test/java/org/springframework/integration/support/MutableMessageTests.java b/spring-integration-core/src/test/java/org/springframework/integration/support/MutableMessageTests.java
index 16b1a24032..dfe923cf5d 100644
--- a/spring-integration-core/src/test/java/org/springframework/integration/support/MutableMessageTests.java
+++ b/spring-integration-core/src/test/java/org/springframework/integration/support/MutableMessageTests.java
@@ -16,6 +16,11 @@
 
 package org.springframework.integration.support;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
 import java.nio.ByteBuffer;
 import java.util.HashMap;
 import java.util.Map;
@@ -26,11 +31,13 @@ import org.junit.Test;
 import org.springframework.messaging.MessageHeaders;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatNoException;
 
 
 /**
  * @author Stuart Williams
  * @author Nathan Kurtyka
+ * @author Mitchell McDonald
  *
  * @since 4.2
  */
@@ -105,4 +112,34 @@ public class MutableMessageTests {
 		assertThat(mutableMessageBytes.getHeaders().getTimestamp()).isEqualTo(timestamp);
 	}
 
+	@Test
+	public void testMessageHeaderIsSerializableAndDeserializableWithNonSerializableValues()
+			throws IOException, ClassNotFoundException {
+
+		String payload = "payload";
+
+		Map<String, Object> headerMap = new HashMap<>();
+		headerMap.put("header1", "serializableValue");
+		headerMap.put("header2", new Object()); // Non-Serializable value
+
+		MutableMessage<String> mutableMessage = new MutableMessage<>(payload, headerMap);
+
+		ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+		ObjectOutputStream outputStream = new ObjectOutputStream(byteArrayOutputStream);
+		outputStream.writeObject(mutableMessage);
+		outputStream.flush();
+
+		ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(byteArrayOutputStream.toByteArray());
+		ObjectInputStream inputStream = new ObjectInputStream(byteArrayInputStream);
+		Object deserializedObject = inputStream.readObject();
+
+		assertThat(deserializedObject).isInstanceOf(MutableMessage.class);
+		MutableMessage<?> deserializedMessage =
+				(MutableMessage<?>) deserializedObject;
+
+		assertThat(deserializedMessage.getHeaders().get("header2")).isNull(); // Non-serializable value removed
+		assertThat(deserializedMessage.getHeaders().get("header1")).isEqualTo("serializableValue");
+		assertThatNoException().isThrownBy(() -> deserializedMessage.getRawHeaders().put("header3", "newValue"));
+	}
+
 }