GH-9546: IMAP: Flag messages after messages have been copied

Fixes: #9546
Issue link: #9546

Sometimes the IMAP connection breaks just after the message flags have been set, but the message has not been copied yet. 
This then leads to the message never being received by (as it has been flagged and the next search will not return it).

* Flag and maybe delete messages after messages have been copied

**Auto-cherry-pick to `6.3.x` & `6.2.x`**

Author: filiphr

Diff for: 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 {

Diff for: 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);
+		}
+	}
+
 }