INT-4276: Selective Header Propagation

JIRA: https://jira.spring.io/browse/INT-4276

The `readOnlyHeaders` integration property allows suppression of certain headers
globally.

Add support for suppressing propagation on individual message handlers.

spring-integration-core/src/main/java/org/springframework/integration/handler/AbstractMessageProducingHandler.java

@@ -16,9 +16,13 @@
 
 package org.springframework.integration.handler;
 
+import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import org.reactivestreams.Publisher;
@@ -37,6 +41,7 @@ import org.springframework.messaging.MessagingException;
 import org.springframework.messaging.core.DestinationResolutionException;
 import org.springframework.messaging.support.ErrorMessage;
 import org.springframework.util.Assert;
+import org.springframework.util.ObjectUtils;
 import org.springframework.util.StringUtils;
 import org.springframework.util.concurrent.ListenableFuture;
 import org.springframework.util.concurrent.ListenableFutureCallback;
@@ -57,6 +62,8 @@ import reactor.core.publisher.Mono;
 public abstract class AbstractMessageProducingHandler extends AbstractMessageHandler
 		implements MessageProducer {
 
+	private final Set<String> notPropagatedHeaders = new HashSet<String>();
+
 	protected final MessagingTemplate messagingTemplate = new MessagingTemplate();
 
 	private volatile MessageChannel outputChannel;
@@ -65,6 +72,8 @@ public abstract class AbstractMessageProducingHandler extends AbstractMessageHan
 
 	private volatile boolean async;
 
+	private boolean selectiveHeaderPropagation;
+
 	/**
 	 * Set the timeout for sending reply Messages.
 	 * @param sendTimeout The send timeout.
@@ -103,6 +112,21 @@ public abstract class AbstractMessageProducingHandler extends AbstractMessageHan
 		return this.async;
 	}
 
+	/**
+	 * Set headers that will NOT be copied from the inbound message if
+	 * {@link #shouldCopyRequestHeaders() shouldCopyRequestHeaaders} is true.
+	 * @param headers the headers to not propagate from the inbound message.
+	 * @since 4.3.10
+	 */
+	public void setNotPropagatedHeaders(String... headers) {
+		if (!ObjectUtils.isEmpty(headers)) {
+			Assert.noNullElements(headers, "null elements are not allowed in 'headers'");
+			this.notPropagatedHeaders.clear();
+			this.notPropagatedHeaders.addAll(Arrays.asList(headers));
+		}
+		this.selectiveHeaderPropagation = this.notPropagatedHeaders.size() > 0;
+	}
+
 	@Override
 	protected void onInit() throws Exception {
 		super.onInit();
@@ -147,7 +171,6 @@ public abstract class AbstractMessageProducingHandler extends AbstractMessageHan
 		return false;
 	}
 
-	@SuppressWarnings("unchecked")
 	protected void produceOutput(Object reply, final Message<?> requestMessage) {
 		final MessageHeaders requestHeaders = requestMessage.getHeaders();
 
@@ -299,7 +322,16 @@ public abstract class AbstractMessageProducingHandler extends AbstractMessageHan
 			builder = this.getMessageBuilderFactory().withPayload(output);
 		}
 		if (this.shouldCopyRequestHeaders()) {
-			builder.copyHeadersIfAbsent(requestHeaders);
+			if (this.selectiveHeaderPropagation) {
+				Map<String, Object> headersToCopy = new HashMap<String, Object>(requestHeaders);
+				for (String header : this.notPropagatedHeaders) {
+					headersToCopy.remove(header);
+				}
+				builder.copyHeadersIfAbsent(headersToCopy);
+			}
+			else {
+				builder.copyHeadersIfAbsent(requestHeaders);
+			}
 		}
 		return builder.build();
 	}

spring-integration-core/src/main/java/org/springframework/integration/handler/support/CollectionArgumentResolver.java

@@ -86,7 +86,7 @@ public class CollectionArgumentResolver extends AbstractExpressionEvaluator
 
 		if (Iterator.class.isAssignableFrom(parameter.getParameterType())) {
 			if (value instanceof Iterable) {
-				return ((Iterable) value).iterator();
+				return ((Iterable<?>) value).iterator();
 			}
 			else {
 				return Collections.singleton(value).iterator();

spring-integration-core/src/test/java/org/springframework/integration/handler/AbstractReplyProducingMessageHandlerTests.java

@@ -17,12 +17,19 @@
 package org.springframework.integration.handler;
 
 import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.CoreMatchers.notNullValue;
+import static org.hamcrest.CoreMatchers.nullValue;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.fail;
-import static org.mockito.Mockito.when;
+import static org.mockito.BDDMockito.given;
+import static org.mockito.BDDMockito.willReturn;
+
+import java.util.Collections;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnitRunner;
 
@@ -30,15 +37,17 @@ import org.springframework.integration.support.MessageBuilder;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
 import org.springframework.messaging.MessagingException;
+import org.springframework.messaging.support.GenericMessage;
 
 /**
  * @author Iwein Fuld
  * @author Gunnar Hillert
+ * @author Gary Russell
  */
 @RunWith(MockitoJUnitRunner.class)
 public class AbstractReplyProducingMessageHandlerTests {
 
-	private AbstractReplyProducingMessageHandler handler = new AbstractReplyProducingMessageHandler() {
+	private final AbstractReplyProducingMessageHandler handler = new AbstractReplyProducingMessageHandler() {
 
 		@Override
 		protected Object handleRequestMessage(Message<?> requestMessage) {
@@ -47,19 +56,19 @@ public class AbstractReplyProducingMessageHandlerTests {
 
 	};
 
-	private Message<?> message = MessageBuilder.withPayload("test").build();
+	private final Message<?> message = MessageBuilder.withPayload("test").build();
 
 	@Mock
-	private MessageChannel channel = null;
+	private final MessageChannel channel = null;
 
 
 	@Test
 	public void errorMessageShouldContainChannelName() {
-		handler.setOutputChannel(channel);
-		when(channel.send(message)).thenReturn(false);
-		when(channel.toString()).thenReturn("testChannel");
+		this.handler.setOutputChannel(this.channel);
+		given(this.channel.send(this.message)).willReturn(false);
+		given(this.channel.toString()).willReturn("testChannel");
 		try {
-			handler.handleMessage(message);
+			this.handler.handleMessage(this.message);
 			fail("Expected a MessagingException");
 		}
 		catch (MessagingException e) {
@@ -67,4 +76,30 @@ public class AbstractReplyProducingMessageHandlerTests {
 		}
 	}
 
+	@Test
+	public void testNotPropagate() {
+		AbstractReplyProducingMessageHandler handler = new AbstractReplyProducingMessageHandler() {
+
+			@Override
+			protected Object handleRequestMessage(Message<?> requestMessage) {
+				return new GenericMessage<String>("world", Collections.singletonMap("bar", "RAB"));
+			}
+
+		};
+		handler.setNotPropagatedHeaders("foo", "bar");
+		handler.setOutputChannel(this.channel);
+		ArgumentCaptor<Message<?>> captor = ArgumentCaptor.forClass(Message.class);
+		willReturn(true).given(this.channel).send(captor.capture());
+		handler.handleMessage(MessageBuilder.withPayload("hello")
+				.setHeader("foo", "FOO")
+				.setHeader("bar", "BAR")
+				.setHeader("baz", "BAZ")
+				.build());
+		Message<?> out = captor.getValue();
+		assertThat(out, notNullValue());
+		assertThat(out.getHeaders().get("foo"), nullValue());
+		assertThat(out.getHeaders().get("bar"), equalTo("RAB"));
+		assertThat(out.getHeaders().get("baz"), equalTo("BAZ"));
+	}
+
 }