From 577c2e87fd0efa89828ab7b5c49357848f8aaee7 Mon Sep 17 00:00:00 2001
From: Artem Bilan <artem.bilan@broadcom.com>
Date: Tue, 24 Sep 2024 09:24:51 -0400
Subject: [PATCH] Fix `AbstractAwsMessageHandler` impls for explicit
 `headerMapper` setting

---
 .../aws/outbound/AbstractAwsMessageHandler.java            | 7 +++++++
 .../integration/aws/outbound/SnsMessageHandler.java        | 4 +++-
 .../integration/aws/outbound/SqsMessageHandler.java        | 4 +++-
 3 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/main/java/org/springframework/integration/aws/outbound/AbstractAwsMessageHandler.java b/src/main/java/org/springframework/integration/aws/outbound/AbstractAwsMessageHandler.java
index 882adc5..f71c80e 100644
--- a/src/main/java/org/springframework/integration/aws/outbound/AbstractAwsMessageHandler.java
+++ b/src/main/java/org/springframework/integration/aws/outbound/AbstractAwsMessageHandler.java
@@ -61,6 +61,8 @@ public abstract class AbstractAwsMessageHandler<H> extends AbstractMessageProduc
 
 	private HeaderMapper<H> headerMapper;
 
+	private boolean headerMapperSet;
+
 	public void setSendTimeout(long sendTimeout) {
 		setSendTimeoutExpression(new ValueExpression<>(sendTimeout));
 	}
@@ -84,6 +86,11 @@ public abstract class AbstractAwsMessageHandler<H> extends AbstractMessageProduc
 	 */
 	public void setHeaderMapper(HeaderMapper<H> headerMapper) {
 		this.headerMapper = headerMapper;
+		this.headerMapperSet = true;
+	}
+
+	protected boolean isHeaderMapperSet() {
+		return this.headerMapperSet;
 	}
 
 	/**
diff --git a/src/main/java/org/springframework/integration/aws/outbound/SnsMessageHandler.java b/src/main/java/org/springframework/integration/aws/outbound/SnsMessageHandler.java
index 7b09866..7880652 100644
--- a/src/main/java/org/springframework/integration/aws/outbound/SnsMessageHandler.java
+++ b/src/main/java/org/springframework/integration/aws/outbound/SnsMessageHandler.java
@@ -187,7 +187,9 @@ public class SnsMessageHandler extends AbstractAwsMessageHandler<Map<String, Mes
 	@Override
 	protected void onInit() {
 		super.onInit();
-		setHeaderMapper(new SnsHeaderMapper());
+		if (!isHeaderMapperSet()) {
+			setHeaderMapper(new SnsHeaderMapper());
+		}
 		TypeLocator typeLocator = getEvaluationContext().getTypeLocator();
 		if (typeLocator instanceof StandardTypeLocator standardTypeLocator) {
 			/*
diff --git a/src/main/java/org/springframework/integration/aws/outbound/SqsMessageHandler.java b/src/main/java/org/springframework/integration/aws/outbound/SqsMessageHandler.java
index fe35da1..31712e9 100644
--- a/src/main/java/org/springframework/integration/aws/outbound/SqsMessageHandler.java
+++ b/src/main/java/org/springframework/integration/aws/outbound/SqsMessageHandler.java
@@ -148,7 +148,9 @@ public class SqsMessageHandler extends AbstractAwsMessageHandler<Map<String, Mes
 	@Override
 	protected void onInit() {
 		super.onInit();
-		setHeaderMapper(new SqsHeaderMapper());
+		if (!isHeaderMapperSet()) {
+			setHeaderMapper(new SqsHeaderMapper());
+		}
 		if (this.messageConverter == null) {
 			this.messageConverter = new GenericMessageConverter(getConversionService());
 		}