From b1663585dc00f80e0b061006484d1325e03d42cb Mon Sep 17 00:00:00 2001
From: Juergen Hoeller <jhoeller@pivotal.io>
Date: Fri, 15 Jul 2016 00:15:46 +0200
Subject: [PATCH] MessageHeaderAccessor properly removes header even in case of
 null value

Issue: SPR-14468
---
 .../support/MessageHeaderAccessor.java        | 13 ++++++----
 .../support/MessageHeaderAccessorTests.java   | 24 ++++++++++++++++---
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/spring-messaging/src/main/java/org/springframework/messaging/support/MessageHeaderAccessor.java b/spring-messaging/src/main/java/org/springframework/messaging/support/MessageHeaderAccessor.java
index 484ab28f967..976c207b413 100644
--- a/spring-messaging/src/main/java/org/springframework/messaging/support/MessageHeaderAccessor.java
+++ b/spring-messaging/src/main/java/org/springframework/messaging/support/MessageHeaderAccessor.java
@@ -306,12 +306,17 @@ public class MessageHeaderAccessor {
 			throw new IllegalArgumentException("'" + name + "' header is read-only");
 		}
 		verifyType(name, value);
-		if (!ObjectUtils.nullSafeEquals(value, getHeader(name))) {
-			this.modified = true;
-			if (value != null) {
+		if (value != null) {
+			// Modify header if necessary
+			if (!ObjectUtils.nullSafeEquals(value, getHeader(name))) {
+				this.modified = true;
 				this.headers.getRawHeaders().put(name, value);
 			}
-			else {
+		}
+		else {
+			// Remove header if available
+			if (this.headers.containsKey(name)) {
+				this.modified = true;
 				this.headers.getRawHeaders().remove(name);
 			}
 		}
diff --git a/spring-messaging/src/test/java/org/springframework/messaging/support/MessageHeaderAccessorTests.java b/spring-messaging/src/test/java/org/springframework/messaging/support/MessageHeaderAccessorTests.java
index a09429fb62b..7d350e8e7a3 100644
--- a/spring-messaging/src/test/java/org/springframework/messaging/support/MessageHeaderAccessorTests.java
+++ b/spring-messaging/src/test/java/org/springframework/messaging/support/MessageHeaderAccessorTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2014 the original author or authors.
+ * Copyright 2002-2016 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.
@@ -18,6 +18,7 @@ package org.springframework.messaging.support;
 
 import java.nio.charset.Charset;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
@@ -40,6 +41,7 @@ import static org.junit.Assert.*;
  *
  * @author Rossen Stoyanchev
  * @author Sebastien Deleuze
+ * @author Juergen Hoeller
  */
 public class MessageHeaderAccessorTests {
 
@@ -89,6 +91,24 @@ public class MessageHeaderAccessorTests {
 		assertEquals("baz", actual.get("bar"));
 	}
 
+	@Test
+	public void testRemoveHeader() {
+		Message<?> message = new GenericMessage<>("payload", Collections.singletonMap("foo", "bar"));
+		MessageHeaderAccessor accessor = new MessageHeaderAccessor(message);
+		accessor.removeHeader("foo");
+		Map<String, Object> headers = accessor.toMap();
+		assertFalse(headers.containsKey("foo"));
+	}
+
+	@Test
+	public void testRemoveHeaderEvenIfNull() {
+		Message<?> message = new GenericMessage<>("payload", Collections.singletonMap("foo", null));
+		MessageHeaderAccessor accessor = new MessageHeaderAccessor(message);
+		accessor.removeHeader("foo");
+		Map<String, Object> headers = accessor.toMap();
+		assertFalse(headers.containsKey("foo"));
+	}
+
 	@Test
 	public void removeHeaders() {
 		Map<String, Object> map = new HashMap<>();
@@ -153,7 +173,6 @@ public class MessageHeaderAccessorTests {
 
 	@Test
 	public void toMap() {
-
 		MessageHeaderAccessor accessor = new MessageHeaderAccessor();
 
 		accessor.setHeader("foo", "bar1");
@@ -380,7 +399,6 @@ public class MessageHeaderAccessorTests {
 	}
 
 
-
 	public static class TestMessageHeaderAccessor extends MessageHeaderAccessor {
 
 		private TestMessageHeaderAccessor() {