Flag and maybe delete messages after messages have been copied

spring-integration-mail/src/main/java/org/springframework/integration/mail/AbstractMailReceiver.java

@@ -68,6 +68,7 @@
  * @author Dominik Simmen
  * @author Yuxin Wang
  * @author Ngoc Nhan
+ * @author Filip Hrisafov
  */
 public abstract class AbstractMailReceiver extends IntegrationObjectSupport implements MailReceiver, DisposableBean {
 
@@ -503,18 +504,27 @@ private Object byteArrayToContent(Map<String, Object> headers, ByteArrayOutputSt
 	}
 
 	private void postProcessFilteredMessages(Message[] filteredMessages) throws MessagingException {
-		setMessageFlags(filteredMessages);
-
-		if (shouldDeleteMessages()) {
-			deleteMessages(filteredMessages);
-		}
 		// Copy messages to cause an eager fetch
+		Message[] messages = filteredMessages;
 		if (this.headerMapper == null && (this.autoCloseFolder || this.simpleContent)) {
+			messages = new Message[filteredMessages.length];
 			for (int i = 0; i < filteredMessages.length; i++) {
-				MimeMessage mimeMessage = new IntegrationMimeMessage((MimeMessage) filteredMessages[i]);
+				Message originalMessage = filteredMessages[i];
+				messages[i] = originalMessage;
+				MimeMessage mimeMessage = new IntegrationMimeMessage((MimeMessage) originalMessage);
 				filteredMessages[i] = mimeMessage;
 			}
 		}
+
+		setMessageFlagsAndMaybeDeleteMessages(messages);
+	}
+
+	private void setMessageFlagsAndMaybeDeleteMessages(Message[] messages) throws MessagingException {
+		setMessageFlags(messages);
+
+		if (shouldDeleteMessages()) {
+			deleteMessages(messages);
+		}
 	}
 
 	private void setMessageFlags(Message[] filteredMessages) throws MessagingException {

spring-integration-mail/src/test/java/org/springframework/integration/mail/ImapMailReceiverTests.java

@@ -17,13 +17,15 @@
 package org.springframework.integration.mail;
 
 import java.io.IOException;
+import java.io.OutputStream;
 import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Properties;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Handler;
@@ -88,6 +90,7 @@
 import org.springframework.util.MimeTypeUtils;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.BDDMockito.given;
@@ -107,6 +110,7 @@
  * @author Artem Bilan
  * @author Alexander Pinske
  * @author Dominik Simmen
+ * @author Filip Hrisafov
  */
 @SpringJUnitConfig
 @ContextConfiguration(
@@ -299,6 +303,11 @@ public void receiveAndMarkAsReadDontDelete() throws Exception {
 
 	private AbstractMailReceiver receiveAndMarkAsReadDontDeleteGuts(AbstractMailReceiver receiver, Message msg1,
 			Message msg2) throws NoSuchFieldException, IllegalAccessException, MessagingException {
+		return receiveAndMarkAsReadDontDeleteGuts(receiver, msg1, msg2, true);
+	}
+
+	private AbstractMailReceiver receiveAndMarkAsReadDontDeleteGuts(AbstractMailReceiver receiver, Message msg1,
+			Message msg2, boolean receive) throws NoSuchFieldException, IllegalAccessException, MessagingException {
 
 		((ImapMailReceiver) receiver).setShouldMarkMessagesAsRead(true);
 		receiver = spy(receiver);
@@ -326,7 +335,9 @@ private AbstractMailReceiver receiveAndMarkAsReadDontDeleteGuts(AbstractMailRece
 		willAnswer(invocation -> messages).given(folder).search(any(SearchTerm.class));
 
 		willAnswer(invocation -> null).given(receiver).fetchMessages(messages);
-		receiver.receive();
+		if (receive) {
+			receiver.receive();
+		}
 		return receiver;
 	}
 
@@ -980,6 +991,31 @@ private void setUpScheduler(ImapMailReceiver mailReceiver, ThreadPoolTaskSchedul
 		mailReceiver.setBeanFactory(bf);
 	}
 
+	@Test
+	public void receiveAndMarkAsReadDontDeleteWithThrowingWhenCopying() throws Exception {
+		AbstractMailReceiver receiver = new ImapMailReceiver();
+		MimeMessage msg1 = spy(GreenMailUtil.newMimeMessage("test1"));
+		MimeMessage greenMailMsg2 = GreenMailUtil.newMimeMessage("test2");
+		TestThrowingMimeMessage msg2 = new TestThrowingMimeMessage(greenMailMsg2);
+		receiver = receiveAndMarkAsReadDontDeleteGuts(receiver, msg1, msg2, false);
+		assertThatThrownBy(receiver::receive)
+				.isInstanceOf(MessagingException.class)
+				.hasMessage("IOException while copying message")
+				.cause()
+				.isInstanceOf(IOException.class)
+				.hasMessage("Simulated exception");
+		assertThat(msg1.getFlags().contains(Flag.SEEN)).isFalse();
+		assertThat(msg2.getFlags().contains(Flag.SEEN)).isFalse();
+		verify(msg1, times(0)).setFlags(Mockito.any(), Mockito.anyBoolean());
+
+		receiver.receive();
+		assertThat(msg1.getFlags().contains(Flag.SEEN)).isTrue();
+		assertThat(msg2.getFlags().contains(Flag.SEEN)).isTrue();
+		// msg2 is marked with the user and seen flags
+		verify(msg1, times(2)).setFlags(Mockito.any(), Mockito.anyBoolean());
+		verify(receiver, times(0)).deleteMessages(Mockito.any());
+	}
+
 	private static class ImapSearchLoggingHandler extends Handler {
 
 		private final List<String> searches = new ArrayList<>();
@@ -1015,4 +1051,21 @@ public void close() throws SecurityException {
 
 	}
 
+	private static class TestThrowingMimeMessage extends MimeMessage {
+
+		protected final AtomicBoolean throwExceptionBeforeWrite = new AtomicBoolean(true);
+
+		private TestThrowingMimeMessage(MimeMessage source) throws MessagingException {
+			super(source);
+		}
+
+		@Override
+		public void writeTo(OutputStream os) throws IOException, MessagingException {
+			if (this.throwExceptionBeforeWrite.getAndSet(false)) {
+				throw new IOException("Simulated exception");
+			}
+			super.writeTo(os);
+		}
+	}
+
 }