From 1d05b013f9215313b6cf8a5d50af7fb107e045a6 Mon Sep 17 00:00:00 2001
From: Oleg Zhurakousky <ozhurakousky@vmware.com>
Date: Fri, 15 Apr 2011 18:08:25 -0400
Subject: [PATCH] INT-1851 added support for mapping JMS properties of the jms
 reply message to SI Message headers

---
 .../integration/jms/JmsOutboundGateway.java         | 13 ++++++++++++-
 .../JmsWithMarshallingMessageConverterTests.java    |  6 ++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/spring-integration-jms/src/main/java/org/springframework/integration/jms/JmsOutboundGateway.java b/spring-integration-jms/src/main/java/org/springframework/integration/jms/JmsOutboundGateway.java
index c7208cfc51..da2c60dd36 100644
--- a/spring-integration-jms/src/main/java/org/springframework/integration/jms/JmsOutboundGateway.java
+++ b/spring-integration-jms/src/main/java/org/springframework/integration/jms/JmsOutboundGateway.java
@@ -16,6 +16,7 @@
 
 package org.springframework.integration.jms;
 
+import java.util.Map;
 import java.util.UUID;
 
 import javax.jms.Connection;
@@ -37,6 +38,7 @@ import org.springframework.integration.MessageChannel;
 import org.springframework.integration.MessageHandlingException;
 import org.springframework.integration.MessageTimeoutException;
 import org.springframework.integration.handler.AbstractReplyProducingMessageHandler;
+import org.springframework.integration.message.GenericMessage;
 import org.springframework.integration.support.MessageBuilder;
 import org.springframework.jms.connection.ConnectionFactoryUtils;
 import org.springframework.jms.support.JmsUtils;
@@ -323,6 +325,7 @@ public class JmsOutboundGateway extends AbstractReplyProducingMessageHandler {
 		}
 	}
 
+	@SuppressWarnings("unchecked")
 	@Override
 	protected Object handleRequestMessage(final Message<?> message) {
 		if (!this.initialized) {
@@ -342,7 +345,15 @@ public class JmsOutboundGateway extends AbstractReplyProducingMessageHandler {
 					logger.debug("converted JMS Message [" + jmsReply + "] to integration Message payload [" + result + "]");
 				}
 			}
-			return result;
+			Map<String, Object> jmsReplyHeaders = (Map<String, Object>) this.headerMapper.toHeaders(jmsReply);
+			Message<?> replyMessage = null;
+			if (result instanceof Message){
+				replyMessage = MessageBuilder.fromMessage((Message<?>) result).copyHeaders(jmsReplyHeaders).build();
+			}
+			else {
+				replyMessage = MessageBuilder.withPayload(result).copyHeaders(jmsReplyHeaders).build();
+			}
+			return replyMessage;
 		}
 		catch (JMSException e) {
 			throw new MessageHandlingException(requestMessage, e);
diff --git a/spring-integration-jms/src/test/java/org/springframework/integration/jms/config/JmsWithMarshallingMessageConverterTests.java b/spring-integration-jms/src/test/java/org/springframework/integration/jms/config/JmsWithMarshallingMessageConverterTests.java
index 85ab5338c7..7682e03a10 100644
--- a/spring-integration-jms/src/test/java/org/springframework/integration/jms/config/JmsWithMarshallingMessageConverterTests.java
+++ b/spring-integration-jms/src/test/java/org/springframework/integration/jms/config/JmsWithMarshallingMessageConverterTests.java
@@ -16,6 +16,7 @@
 
 package org.springframework.integration.jms.config;
 
+import static junit.framework.Assert.assertNotNull;
 import static org.junit.Assert.assertEquals;
 
 import java.io.IOException;
@@ -32,6 +33,7 @@ import org.springframework.context.ApplicationContext;
 import org.springframework.context.support.ClassPathXmlApplicationContext;
 import org.springframework.integration.Message;
 import org.springframework.integration.MessageChannel;
+import org.springframework.integration.MessageHeaders;
 import org.springframework.integration.core.PollableChannel;
 import org.springframework.integration.message.GenericMessage;
 import org.springframework.oxm.Marshaller;
@@ -53,6 +55,10 @@ public class JmsWithMarshallingMessageConverterTests {
 		PollableChannel output = ac.getBean("output", PollableChannel.class);
 		input.send(new GenericMessage<String>("hello"));
 		Message<String> replyMessage = (Message<String>) output.receive();
+		MessageHeaders headers = replyMessage.getHeaders();
+		// check for couple of JMS headers, make sure they are present
+		assertNotNull(headers.get("jms_redelivered"));
+		assertNotNull(headers.get("jms_correlationId"));
 		assertEquals("HELLO", replyMessage.getPayload());
 	}