Flag and maybe delete messages after messages have been copied

Author: filiphr

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

@@ -503,17 +503,26 @@ 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
 		if (this.headerMapper == null && (this.autoCloseFolder || this.simpleContent)) {
+			Message[] originalMessages = new Message[filteredMessages.length];
 			for (int i = 0; i < filteredMessages.length; i++) {
-				MimeMessage mimeMessage = new IntegrationMimeMessage((MimeMessage) filteredMessages[i]);
+				Message originalMessage = filteredMessages[i];
+				originalMessages[i] = originalMessage;
+				MimeMessage mimeMessage = new IntegrationMimeMessage((MimeMessage) originalMessage);
 				filteredMessages[i] = mimeMessage;
 			}
+			setMessageFlagsAndMaybeDeleteMessages(originalMessages);
+		} else {
+			setMessageFlagsAndMaybeDeleteMessages(filteredMessages);
+		}
+	}
+
+	private void setMessageFlagsAndMaybeDeleteMessages(Message[] messages) throws MessagingException {
+		setMessageFlags(messages);
+
+		if (shouldDeleteMessages()) {
+			deleteMessages(messages);
 		}
 	}
 

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

@@ -17,6 +17,7 @@
 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;
@@ -88,6 +89,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;
@@ -299,6 +301,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 +333,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 +989,28 @@ private void setUpScheduler(ImapMailReceiver mailReceiver, ThreadPoolTaskSchedul
 		mailReceiver.setBeanFactory(bf);
 	}
 
+	@Test
+	public void receiveAndMarkAsReadDontDeleteWithThrowingWhenCopying() throws Exception {
+		AbstractMailReceiver receiver = new ImapMailReceiver();
+		MimeMessage msg1 = GreenMailUtil.newMimeMessage("test1");
+		MimeMessage greenMailMsg2 = GreenMailUtil.newMimeMessage("test2");
+		TestThrowingMimeMessage msg2 = new TestThrowingMimeMessage(greenMailMsg2, 1);
+		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();
+
+		receiver.receive();
+		assertThat(msg1.getFlags().contains(Flag.SEEN)).isTrue();
+		assertThat(msg2.getFlags().contains(Flag.SEEN)).isTrue();
+		verify(receiver, times(0)).deleteMessages(Mockito.any());
+	}
+
 	private static class ImapSearchLoggingHandler extends Handler {
 
 		private final List<String> searches = new ArrayList<>();
@@ -1015,4 +1046,22 @@ public void close() throws SecurityException {
 
 	}
 
+	private static class TestThrowingMimeMessage extends MimeMessage {
+
+		protected final AtomicInteger exceptionsBeforeWrite;
+
+		public TestThrowingMimeMessage(MimeMessage source, int exceptionsBeforeWrite) throws MessagingException {
+			super(source);
+			this.exceptionsBeforeWrite = new AtomicInteger(exceptionsBeforeWrite);
+		}
+
+		@Override
+		public void writeTo(OutputStream os) throws IOException, MessagingException {
+			if (this.exceptionsBeforeWrite.decrementAndGet() >= 0) {
+				throw new IOException("Simulated exception");
+			}
+			super.writeTo(os);
+		}
+	}
+
 }