From 021ee6faf21bd553e742640d6cc3a58b8cdc798d Mon Sep 17 00:00:00 2001
From: Dustin Rodrigues <dust.rod@gmail.com>
Date: Sat, 16 Jun 2018 14:01:00 -0400
Subject: [PATCH] NIFI-5317 - support non-ASCII X-Mailer PutEmail header

Signed-off-by: Pierre Villard <pierre.villard.fr@gmail.com>

This closes #2800.
---
 .../nifi/processors/standard/PutEmail.java      | 17 +++++++++++------
 .../nifi/processors/standard/TestPutEmail.java  |  4 ++--
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/PutEmail.java b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/PutEmail.java
index b80f401709..8fdcc14e8e 100644
--- a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/PutEmail.java
+++ b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/main/java/org/apache/nifi/processors/standard/PutEmail.java
@@ -320,6 +320,15 @@ public class PutEmail extends AbstractProcessor {
         this.attributeNamePattern = attributeNameRegex == null ? null : Pattern.compile(attributeNameRegex);
     }
 
+    private void setMessageHeader(final String header, final String value, final Message message) throws MessagingException {
+        final ComponentLog logger = getLogger();
+        try {
+            message.setHeader(header, MimeUtility.encodeText(value));
+        } catch (UnsupportedEncodingException e){
+            logger.warn("Unable to add header {} with value {} due to encoding exception", new Object[]{header, value});
+        }
+    }
+
     @Override
     public void onTrigger(final ProcessContext context, final ProcessSession session) {
         final FlowFile flowFile = session.get();
@@ -343,15 +352,11 @@ public class PutEmail extends AbstractProcessor {
             if (attributeNamePattern != null) {
                 for (final Map.Entry<String, String> entry : flowFile.getAttributes().entrySet()) {
                     if (attributeNamePattern.matcher(entry.getKey()).matches()) {
-                        try {
-                            message.setHeader(entry.getKey(), MimeUtility.encodeText(entry.getValue()));
-                        } catch (UnsupportedEncodingException e){
-                            logger.warn("Unable to add header value {} due to encoding exception", new Object[]{entry.getValue()});
-                        }
+                        this.setMessageHeader(entry.getKey(), entry.getValue(), message);
                     }
                 }
             }
-            message.setHeader("X-Mailer", context.getProperty(HEADER_XMAILER).evaluateAttributeExpressions(flowFile).getValue());
+            this.setMessageHeader("X-Mailer", context.getProperty(HEADER_XMAILER).evaluateAttributeExpressions(flowFile).getValue(), message);
             message.setSubject(context.getProperty(SUBJECT).evaluateAttributeExpressions(flowFile).getValue());
 
             String messageText = getMessage(flowFile, context, session);
diff --git a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestPutEmail.java b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestPutEmail.java
index 2509868e13..7df04ba61d 100644
--- a/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestPutEmail.java
+++ b/nifi-nar-bundles/nifi-standard-bundle/nifi-standard-processors/src/test/java/org/apache/nifi/processors/standard/TestPutEmail.java
@@ -137,7 +137,7 @@ public class TestPutEmail {
     public void testOutgoingMessageWithOptionalProperties() throws Exception {
         // verifies that optional attributes are set on the outgoing Message correctly
         runner.setProperty(PutEmail.SMTP_HOSTNAME, "smtp-host");
-        runner.setProperty(PutEmail.HEADER_XMAILER, "TestingNiFi");
+        runner.setProperty(PutEmail.HEADER_XMAILER, "TestingNíFiNonASCII");
         runner.setProperty(PutEmail.FROM, "${from}");
         runner.setProperty(PutEmail.MESSAGE, "${message}");
         runner.setProperty(PutEmail.TO, "${to}");
@@ -164,7 +164,7 @@ public class TestPutEmail {
         assertEquals("Expected a single message to be sent", 1, processor.getMessages().size());
         Message message = processor.getMessages().get(0);
         assertEquals("\"test@apache.org\" <NiFi>", message.getFrom()[0].toString());
-        assertEquals("X-Mailer Header", "TestingNiFi", message.getHeader("X-Mailer")[0]);
+        assertEquals("X-Mailer Header", "TestingNíFiNonASCII", MimeUtility.decodeText(message.getHeader("X-Mailer")[0]));
         assertEquals("the message body", message.getContent());
         assertEquals(1, message.getRecipients(RecipientType.TO).length);
         assertEquals("to@apache.org", message.getRecipients(RecipientType.TO)[0].toString());