NIFI-1899 reworking bi-directional Email processing in ListenSMTP

This closes #483

nifi-nar-bundles/nifi-email-bundle/nifi-email-processors/src/main/java/org/apache/nifi/processors/email/ListenSMTP.java

@@ -20,7 +20,7 @@ import javax.net.ssl.SSLContext;
 import javax.net.ssl.SSLSocket;
 import javax.net.ssl.SSLSocketFactory;
 import java.io.IOException;
-import java.io.OutputStream;
+import java.io.InputStream;
 import java.net.InetSocketAddress;
 import java.net.Socket;
 import java.util.ArrayList;
@@ -44,7 +44,6 @@ import org.apache.nifi.annotation.lifecycle.OnStopped;
 import org.apache.nifi.annotation.lifecycle.OnUnscheduled;
 import org.apache.nifi.flowfile.attributes.CoreAttributes;
 import org.apache.nifi.processor.DataUnit;
-import org.apache.nifi.stream.io.ByteArrayOutputStream;
 
 import org.apache.nifi.annotation.lifecycle.OnScheduled;
 import org.apache.nifi.components.PropertyDescriptor;
@@ -64,7 +63,6 @@ import org.apache.nifi.processor.exception.ProcessException;
 import org.apache.nifi.components.ValidationContext;
 import org.apache.nifi.components.ValidationResult;
 import org.apache.nifi.flowfile.FlowFile;
-import org.apache.nifi.processor.io.OutputStreamCallback;
 import org.apache.nifi.ssl.SSLContextService;
 
 import org.apache.nifi.processors.email.smtp.event.SmtpEvent;
@@ -356,27 +354,46 @@ public class ListenSMTP extends AbstractProcessor {
         while (!incomingMessages.isEmpty()) {
             SmtpEvent message = incomingMessages.poll();
 
-
             if (message == null) {
                 return;
             }
 
             synchronized (message) {
+                if (resultCodeSetAndIsError(message)) {
+                    SMTPResultCode resultCode = SMTPResultCode.fromCode(message.getReturnCode());
+                    getLogger().warn("Message failed before onTrigger processing message was: " + resultCode.getLogMessage());
+                    continue;
+                }
 
+                try {
                     FlowFile flowfile = session.create();
 
                     if (message.getMessageData() != null) {
-                    ByteArrayOutputStream messageData = message.getMessageData();
+                        flowfile = session.write(flowfile, out -> {
-                    flowfile = session.write(flowfile, new OutputStreamCallback() {
+                            InputStream inputStream = message.getMessageData();
+                            byte [] buffer = new byte[1024];
 
-                        // Write the messageData to flowfile content
+                            int rd;
-                        @Override
+                            long totalBytesRead =0;
-                        public void process(OutputStream out) throws IOException {
+
-                            out.write(messageData.toByteArray());
+                            while ((rd = inputStream.read(buffer, 0, buffer.length)) != -1 ) {
+                                totalBytesRead += rd;
+                                if (totalBytesRead > server.getMaxMessageSize() ) {
+                                    message.setReturnCode(500);
+                                    message.setProcessed();
+                                    break;
                                 }
+                                out.write(buffer, 0, rd);
+                            }
+                            out.flush();
                         });
+                    } else {
+                        getLogger().debug("Message body was null");
+                        message.setReturnCode(SMTPResultCode.UNKNOWN_ERROR_CODE.getCode());
+                        message.setProcessed();
                     }
 
+                    if (!message.getProcessed()) {
                         HashMap<String, String> attributes = new HashMap<>();
                         // Gather message attributes
                         attributes.put(SMTP_HELO, message.getHelo());
@@ -402,41 +419,60 @@ public class ListenSMTP extends AbstractProcessor {
                         flowfile = session.putAllAttributes(flowfile, attributes);
                         session.getProvenanceReporter().receive(flowfile, "smtp://" + SMTP_HOSTNAME + ":" + SMTP_PORT + "/");
                         session.transfer(flowfile, REL_SUCCESS);
+
                         getLogger().info("Transferring {} to success", new Object[]{flowfile});
+                    }
+                } catch (Exception e) {
+                    message.setProcessed();
+                    message.setReturnCode(SMTPResultCode.UNEXPECTED_ERROR.getCode());
+                }
+
+                // Check to see if it failed when creating the FlowFile
+                if (resultCodeSetAndIsError(message)) {
+                    session.rollback();
+                    SMTPResultCode resultCode = SMTPResultCode.fromCode(message.getReturnCode());
+                    getLogger().warn("Failed to received message due to: " + resultCode.getLogMessage());
+                    message.notifyAll();
+                    continue;
+                }
 
                 // Finished processing,
                 message.setProcessed();
 
-                // update the latch so data() can process the rest of the method
+                // notify on the message so data() can process the rest of the method
-                message.updateProcessedLatch();
+                message.notifyAll();
 
-                // End of synchronized block
+                // Wait for data() to tell sender we received the message and double check we didn't timeout
+                final long serverTimeout = context.getProperty(SMTP_TIMEOUT).asTimePeriod(TimeUnit.MILLISECONDS);
+                try {
+                    message.wait(serverTimeout);
+                } catch (InterruptedException e) {
+                    getLogger().info("Interrupted while waiting for Message Handler to acknowledge message.");
                 }
 
-            // Wait for SMTPMessageHandler data() and done() to complete
+                // Check to see if the sender was correctly notified
-            // their side of the work (i.e. acknowledgement)
+                if (resultCodeSetAndIsError(message)) {
-            while (!message.getAcknowledged()) {
-                // Busy wait
-                }
-
-            // Lock one last time
-            synchronized (message) {
                     SMTPResultCode resultCode = SMTPResultCode.fromCode(message.getReturnCode());
-                switch (resultCode) {
-                    case UNEXPECTED_ERROR:
-                    case TIMEOUT_ERROR:
                     session.rollback();
-                        getLogger().warn(resultCode.getLogMessage());
+                    getLogger().warn("Failed to received message due to: " + resultCode.getLogMessage());
-                    case SUCCESS:
+                } else {
-                        getLogger().info(resultCode.getLogMessage());
+                    // Need to commit because if we didn't and a following message needed to be rolled back, this message would be too, causing data loss.
-                        break;
+                    session.commit();
-                    default:
-                        getLogger().error(resultCode.getLogMessage());
                 }
             }
         }
     }
 
+    private boolean resultCodeSetAndIsError(SmtpEvent message){
+        if (message.getReturnCode() != null ) {
+            SMTPResultCode resultCode = SMTPResultCode.fromCode(message.getReturnCode());
+            if (resultCode.isError()) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     // Same old... same old... used for testing to access the random port that was selected
     protected int getPort() {
         return server == null ? 0 : server.getPort();

nifi-nar-bundles/nifi-email-bundle/nifi-email-processors/src/main/java/org/apache/nifi/processors/email/smtp/event/SmtpEvent.java

@@ -18,14 +18,13 @@
 package org.apache.nifi.processors.email.smtp.event;
 
 
-import org.apache.nifi.stream.io.ByteArrayOutputStream;
 
+import java.io.InputStream;
 import java.security.cert.X509Certificate;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -38,19 +37,15 @@ public class SmtpEvent{
     private final String helo;
     private final String from;
     private final String to;
-    private final ByteArrayOutputStream messageData;
+    private final InputStream messageData;
     private List<Map<String, String>> certificatesDetails;
     private AtomicBoolean processed = new AtomicBoolean(false);
     private AtomicBoolean acknowledged = new AtomicBoolean(false);
     private AtomicInteger returnCode = new AtomicInteger();
-    private CountDownLatch processedLatch;
 
     public SmtpEvent(
             final String remoteIP, final String helo, final String from, final String to, final X509Certificate[] certificates,
-            final ByteArrayOutputStream messageData,
+            final InputStream messageData) {
-            CountDownLatch processedLatch) {
-
-        this.processedLatch = processedLatch;
 
         this.remoteIP = remoteIP;
         this.helo = helo;
@@ -86,7 +81,7 @@ public class SmtpEvent{
         return helo;
     }
 
-    public synchronized ByteArrayOutputStream getMessageData() {
+    public synchronized InputStream getMessageData() {
         return messageData;
     }
 
@@ -118,15 +113,11 @@ public class SmtpEvent{
         return this.acknowledged.get();
     }
 
-    public synchronized void updateProcessedLatch() {
-        this.processedLatch.countDown();
-    }
-
     public synchronized void setReturnCode(int code) {
         this.returnCode.set(code);
     }
 
-    public synchronized int getReturnCode() {
+    public synchronized Integer getReturnCode() {
         return this.returnCode.get();
     }
 

nifi-nar-bundles/nifi-email-bundle/nifi-email-processors/src/main/java/org/apache/nifi/processors/email/smtp/handler/SMTPMessageHandlerFactory.java

@@ -20,12 +20,10 @@ package org.apache.nifi.processors.email.smtp.handler;
 import java.io.IOException;
 import java.io.InputStream;
 import java.security.cert.X509Certificate;
-import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 
 import org.apache.nifi.logging.ComponentLog;
-import org.apache.nifi.stream.io.ByteArrayOutputStream;
 import org.apache.nifi.util.StopWatch;
 import org.subethamail.smtp.DropConnectionException;
 import org.subethamail.smtp.MessageContext;
@@ -33,7 +31,6 @@ import org.subethamail.smtp.MessageHandler;
 import org.subethamail.smtp.MessageHandlerFactory;
 import org.subethamail.smtp.RejectException;
 import org.subethamail.smtp.TooMuchDataException;
-import org.subethamail.smtp.server.SMTPServer;
 
 import org.apache.nifi.processors.email.smtp.event.SmtpEvent;
 
@@ -57,13 +54,9 @@ public class SMTPMessageHandlerFactory implements MessageHandlerFactory {
         final MessageContext messageContext;
         String from;
         String recipient;
-        ByteArrayOutputStream messageData;
-
-        private CountDownLatch latch;
 
         public Handler(MessageContext messageContext, LinkedBlockingQueue<SmtpEvent> incomingMessages, ComponentLog logger){
             this.messageContext = messageContext;
-            this.latch =  new CountDownLatch(1);
         }
 
         @Override
@@ -82,31 +75,9 @@ public class SMTPMessageHandlerFactory implements MessageHandlerFactory {
         public void data(InputStream inputStream) throws RejectException, TooMuchDataException, IOException {
             // Start counting the timer...
             StopWatch watch = new StopWatch(true);
-
             long elapsed;
-
-            SMTPServer server = messageContext.getSMTPServer();
-
             final long serverTimeout = TimeUnit.MILLISECONDS.convert(messageContext.getSMTPServer().getConnectionTimeout(), TimeUnit.MILLISECONDS);
 
-            this.messageData = new ByteArrayOutputStream();
-
-            byte [] buffer = new byte[1024];
-
-            int rd;
-
-            while ((rd = inputStream.read(buffer, 0, buffer.length)) != -1 ) {
-                messageData.write(buffer, 0, rd);
-                if (messageData.getBufferLength() > server.getMaxMessageSize() ) {
-                    // NOTE: Setting processed at this stage is not desirable as message object will only be created
-                    // if this test (i.e. message size) passes.
-                    final SMTPResultCode returnCode = SMTPResultCode.fromCode(500);
-                    logger.warn(returnCode.getLogMessage());
-                    throw new TooMuchDataException(returnCode.getErrorMessage());
-                }
-            }
-            messageData.flush();
-
             X509Certificate[] certificates = new X509Certificate[]{};
 
             final String remoteIP = messageContext.getRemoteAddress().toString();
@@ -116,8 +87,9 @@ public class SMTPMessageHandlerFactory implements MessageHandlerFactory {
                 certificates = (X509Certificate[]) messageContext.getTlsPeerCertificates().clone();
             }
 
-            SmtpEvent message = new SmtpEvent(remoteIP, helo, from, recipient, certificates, messageData, latch);
+            SmtpEvent message = new SmtpEvent(remoteIP, helo, from, recipient, certificates, inputStream);
 
+            synchronized (message) {
                 // / Try to queue the message back to the NiFi session
                 try {
                     elapsed = watch.getElapsed(TimeUnit.MILLISECONDS);
@@ -126,7 +98,7 @@ public class SMTPMessageHandlerFactory implements MessageHandlerFactory {
                     final SMTPResultCode returnCode = SMTPResultCode.fromCode(421);
                     logger.trace(returnCode.getLogMessage());
 
-                // NOTE: Setting processed at this stage is redundant as this catch deals with the inability of
+                    // NOTE: Setting acknowledged at this stage is redundant as this catch deals with the inability of
                     // adding message to the processing queue. Yet, for the sake of consistency the message is
                     // updated nonetheless
                     message.setReturnCode(returnCode.getCode());
@@ -139,10 +111,10 @@ public class SMTPMessageHandlerFactory implements MessageHandlerFactory {
                 // given back to the SMTP client
                 elapsed = watch.getElapsed(TimeUnit.MILLISECONDS);
                 try {
-                latch.await(serverTimeout - elapsed, TimeUnit.MILLISECONDS);
+                    message.wait(serverTimeout - elapsed);
                 } catch (InterruptedException e) {
-                // Latch open unexpectedly. Will return error and requestonTrigger to rollback
+                    // Interrupted while waiting for the message to process. Will return error and request onTrigger to rollback
-                logger.trace("Latch opened unexpectedly and processor indicates data wasn't processed. Returned error to SMTP client as precautionary measure");
+                    logger.trace("Interrupted while waiting for processor to process data. Returned error to SMTP client as precautionary measure");
                     incomingMessages.remove(message);
 
                     // Set the final values so onTrigger can figure out what happened to message
@@ -154,29 +126,40 @@ public class SMTPMessageHandlerFactory implements MessageHandlerFactory {
                     throw new DropConnectionException(returnCode.getCode(), returnCode.getErrorMessage());
                 }
 
-            // Remove the message from the queue.
+                // Check if message is processed
+                if (!message.getProcessed()) {
                     incomingMessages.remove(message);
-            // Check if message is processed and if yes, check if it was received on time and wraps it up.
-            elapsed = watch.getElapsed(TimeUnit.MILLISECONDS);
-            if (!message.getProcessed() ||  (elapsed >= serverTimeout)) {
                     final SMTPResultCode returnCode = SMTPResultCode.fromCode(451);
-                logger.trace("Did not receive the onTrigger response within the acceptable timeframes. Data duplication may have occurred.");
+                    logger.trace("Did not receive the onTrigger response within the acceptable timeframe.");
 
                     // Set the final values so onTrigger can figure out what happened to message
                     message.setReturnCode(returnCode.getCode());
                     message.setAcknowledged();
                     throw new DropConnectionException(returnCode.getCode(), returnCode.getErrorMessage());
-            }
+                } else if(message.getReturnCode() != null) {
+                    // No need to check if over server timeout because we already processed the data. Might as well use the status code returned by onTrigger.
+                    final SMTPResultCode returnCode = SMTPResultCode.fromCode(message.getReturnCode());
 
+                    if(returnCode.isError()){
+                        message.setAcknowledged();
+                        throw new DropConnectionException(returnCode.getCode(), returnCode.getErrorMessage());
+                    }
+                } else {
+                    // onTrigger successfully processed the data.
+                    // No need to check if over server timeout because we already processed the data. Might as well finalize it.
                     // Set the final values so onTrigger can figure out what happened to message
                     message.setReturnCode(250);
                     message.setAcknowledged();
+                }
                 // Exit, allowing Handler to acknowledge the message
+                message.notifyAll();
+            }
         }
 
         @Override
         public void done() {
             logger.trace("Called the last method of message handler. Exiting");
+            // Notifying the ontrigger that the message was handled.
         }
     }
 }

nifi-nar-bundles/nifi-email-bundle/nifi-email-processors/src/main/java/org/apache/nifi/processors/email/smtp/handler/SMTPResultCode.java

@@ -79,5 +79,15 @@ public enum SMTPResultCode {
         return (smtpResultCode == null) ? UNKNOWN_ERROR_CODE : smtpResultCode;
     }
 
-
+    public boolean isError(){
+        switch (this) {
+            case MESSAGE_TOO_LARGE:
+            case UNEXPECTED_ERROR:
+            case QUEUE_ERROR:
+            case TIMEOUT_ERROR:
+                return true;
+            default:
+                return false;
+        }
+    }
 }

nifi-nar-bundles/nifi-email-bundle/nifi-email-processors/src/test/java/org/apache/nifi/processors/email/TestListenSMTP.java

@@ -32,6 +32,7 @@ import org.junit.Assert;
 import org.junit.Test;
 
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 public class TestListenSMTP {
 
@@ -251,7 +252,7 @@ public class TestListenSMTP {
         listenSmtp.startShutdown();
     }
 
-    @Test(timeout=15000, expected=EmailException.class)
+    @Test(timeout=15000)
     public void emailTooLarge() throws Exception {
         ListenSMTP listenSmtp = new ListenSMTP();
         final TestRunner runner = TestRunners.newTestRunner(listenSmtp);
@@ -272,7 +273,14 @@ public class TestListenSMTP {
         listenSmtp.initializeSMTPServer(context);
 
         final int port = listenSmtp.getPort();
+        AtomicBoolean finished = new AtomicBoolean(false);;
+        AtomicBoolean failed = new AtomicBoolean(false);
 
+        try {
+            final Thread clientThread = new Thread(new Runnable() {
+                @Override
+                public void run() {
+                    try {
                         Email email = new SimpleEmail();
                         email.setHostName("127.0.0.1");
                         email.setSmtpPort(port);
@@ -283,20 +291,29 @@ public class TestListenSMTP {
                         email.addTo("bob@nifi.apache.org");
                         email.send();
 
-        Thread.sleep(100);
+                    } catch (final EmailException t) {
-
+                        failed.set(true);
+                    }
+                    finished.set(true);
+                }
+            });
+            clientThread.start();
 
+            while (!finished.get()) {
                 // process the request.
                 listenSmtp.onTrigger(context, processSessionFactory);
+                Thread.sleep(10);
+            }
+            clientThread.stop();
+
+            Assert.assertTrue("Sending email succeeded when it should have failed", failed.get());
 
             runner.assertTransferCount(ListenSMTP.REL_SUCCESS, 0);
-        runner.assertQueueEmpty();
 
-        try {
+            runner.assertQueueEmpty();
+        } finally {
+            // shut down the server
             listenSmtp.startShutdown();
-        } catch (InterruptedException e) {
-                e.printStackTrace();
-                Assert.assertFalse(e.toString(), true);
         }
     }
 }