From 1eee339c152a3de2f2a37525899be28bb0bfaf0a Mon Sep 17 00:00:00 2001
From: Gary Russell <grussell@gopivotal.com>
Date: Wed, 2 Apr 2014 12:49:58 -0400
Subject: [PATCH] Relax constraints in MessageHeaders for subclasses

Allow a subclass to modify MessageHeaders as well as override the
strategy for or even skip having `ID` and `TIMESTAMP` headers.

Issue: SPR-11468
---
 .../messaging/MessageHeaders.java             | 19 ++++++++++++++++--
 .../messaging/MessageHeadersTests.java        | 20 ++++++++++++++++++-
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/spring-messaging/src/main/java/org/springframework/messaging/MessageHeaders.java b/spring-messaging/src/main/java/org/springframework/messaging/MessageHeaders.java
index 0f15080521f..02a73c18ff7 100644
--- a/spring-messaging/src/main/java/org/springframework/messaging/MessageHeaders.java
+++ b/spring-messaging/src/main/java/org/springframework/messaging/MessageHeaders.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2013 the original author or authors.
+ * Copyright 2002-2014 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.
@@ -42,6 +42,7 @@ import org.springframework.util.IdGenerator;
  * <b>IMPORTANT</b>: This class is immutable. Any mutating operation such as
  * {@code put(..)}, {@code putAll(..)} and others will throw
  * {@link UnsupportedOperationException}.
+ * <p>Subclasses do have access to the raw headers, however, via {@link #getRawHeaders()}.
  * <p>
  * One way to create message headers is to use the
  * {@link org.springframework.messaging.support.MessageBuilder MessageBuilder}:
@@ -66,7 +67,7 @@ import org.springframework.util.IdGenerator;
  * @see org.springframework.messaging.support.MessageBuilder
  * @see org.springframework.messaging.support.MessageHeaderAccessor
  */
-public final class MessageHeaders implements Map<String, Object>, Serializable {
+public class MessageHeaders implements Map<String, Object>, Serializable {
 
 	private static final long serialVersionUID = -4615750558355702881L;
 
@@ -95,13 +96,27 @@ public final class MessageHeaders implements Map<String, Object>, Serializable {
 
 	private final Map<String, Object> headers;
 
+	/**
+	 * Constructs a minimal {@link MessageHeaders} with zero headers.
+	 */
+	protected MessageHeaders() {
+		this.headers = new HashMap<String, Object>();
+	}
 
+	/**
+	 * Consructs a {@link MessageHeaders} from the headers map; adding (or
+	 * overwriting) the {@link #ID} and {@link #TIMESTAMP} headers.
+	 * @param headers The map.
+	 */
 	public MessageHeaders(Map<String, Object> headers) {
 		this.headers = (headers != null) ? new HashMap<String, Object>(headers) : new HashMap<String, Object>();
 		this.headers.put(ID, ((idGenerator != null) ? idGenerator : defaultIdGenerator).generateId());
 		this.headers.put(TIMESTAMP, System.currentTimeMillis());
 	}
 
+	protected Map<String, Object> getRawHeaders() {
+		return this.headers;
+	}
 
 	public UUID getId() {
 		return this.get(ID, UUID.class);
diff --git a/spring-messaging/src/test/java/org/springframework/messaging/MessageHeadersTests.java b/spring-messaging/src/test/java/org/springframework/messaging/MessageHeadersTests.java
index 7b7446caa8c..0c88c39f143 100644
--- a/spring-messaging/src/test/java/org/springframework/messaging/MessageHeadersTests.java
+++ b/spring-messaging/src/test/java/org/springframework/messaging/MessageHeadersTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2013 the original author or authors.
+ * Copyright 2002-2014 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.
@@ -23,6 +23,8 @@ import java.io.ObjectOutputStream;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.atomic.AtomicLong;
 
 import org.junit.Test;
 
@@ -32,6 +34,7 @@ import static org.junit.Assert.*;
  * Test fixture for {@link MessageHeaders}.
  *
  * @author Rossen Stoyanchev
+ * @author Gary Russell
  */
 public class MessageHeadersTests {
 
@@ -138,6 +141,21 @@ public class MessageHeadersTests {
 		assertNull(output.get("address"));
 	}
 
+	@Test
+	public void subclass() {
+		final AtomicLong id = new AtomicLong();
+		@SuppressWarnings("serial")
+		class MyMH extends MessageHeaders {
+
+			public MyMH() {
+				this.getRawHeaders().put(ID, new UUID(0, id.incrementAndGet()));
+			}
+
+		}
+		MessageHeaders headers = new MyMH();
+		assertEquals("00000000-0000-0000-0000-000000000001", headers.getId().toString());
+		assertEquals(1, headers.size());
+	}
 
 	private static Object serializeAndDeserialize(Object object) throws Exception {
 		ByteArrayOutputStream baos = new ByteArrayOutputStream();