[ENHANCEMENT] Ensure the MDC context is carried over upon IMAP parsin…

…g errors This would ease linking and identifying IMAP parsing issues
apache · Nov 12, 2024 · 105a0ac · 105a0ac
1 parent 61b6018
commit 105a0ac
Show file tree

Hide file tree

Showing 7 changed files with 49 additions and 23 deletions.
diff --git a/protocols/imap/src/main/java/org/apache/james/imap/api/process/ImapSession.java b/protocols/imap/src/main/java/org/apache/james/imap/api/process/ImapSession.java
@@ -19,6 +19,8 @@
 
 package org.apache.james.imap.api.process;
 
+import java.io.Closeable;
+import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.time.Duration;
 import java.util.Objects;
@@ -32,6 +34,7 @@
 import org.apache.james.mailbox.MailboxSession;
 import org.apache.james.protocols.api.CommandDetectionSession;
 import org.apache.james.protocols.api.OidcSASLConfiguration;
+import org.apache.james.util.MDCBuilder;
 
 import reactor.core.publisher.Mono;
 
@@ -43,6 +46,8 @@
  * @version $Revision: 109034 $
  */
 public interface ImapSession extends CommandDetectionSession {
+    String MDC_KEY = "bound_MDC";
+
     class SessionId {
         private static final RandomStringGenerator RANDOM_STRING_GENERATOR = new RandomStringGenerator.Builder()
             .withinRange('a', 'z')
@@ -256,6 +261,21 @@ default boolean backpressureNeeded(Runnable restoreBackpressure) {
 
     boolean supportsOAuth();
 
+    default void withMDC(Runnable runnable) {
+        try (Closeable c = mdc().build()) {
+            runnable.run();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    default MDCBuilder mdc() {
+        if (getAttribute(MDC_KEY) instanceof MDCBuilder mdcBuilder) {
+            return mdcBuilder;
+        }
+        return MDCBuilder.create();
+    }
+
     /**
      * Return the {@link InetSocketAddress} of the remote peer
      */

diff --git a/...ocols/imap/src/main/java/org/apache/james/imap/decode/base/AbstractImapCommandParser.java b/...ocols/imap/src/main/java/org/apache/james/imap/decode/base/AbstractImapCommandParser.java
@@ -71,7 +71,7 @@ public final ImapMessage parse(ImapRequestLineReader request, Tag tag, ImapSessi
             try {
                 return decode(request, tag, session);
             } catch (DecodingException e) {
-                LOGGER.debug("Cannot parse protocol ", e);
+                session.withMDC(() -> LOGGER.info("Cannot parse protocol ", e));
                 return statusResponseFactory.taggedBad(tag, command, e.getKey());
             }
         }

diff --git a/protocols/imap/src/main/java/org/apache/james/imap/decode/main/DefaultImapDecoder.java b/protocols/imap/src/main/java/org/apache/james/imap/decode/main/DefaultImapDecoder.java
@@ -70,7 +70,7 @@ public ImapMessage decode(ImapRequestLineReader request, ImapSession session) {
             Tag tag = request.tag();
             return decodeCommandTagged(request, tag, session);
         } catch (DecodingException e) {
-            LOGGER.debug("Cannot parse tag", e);
+            session.withMDC(() -> LOGGER.debug("Cannot parse tag", e));
             return unknownCommand(null, session);
         }
     }
@@ -80,7 +80,7 @@ private ImapMessage decodeCommandTagged(ImapRequestLineReader request, Tag tag,
             String commandName = request.atom();
             return decodeCommandNamed(request, tag, commandName, session);
         } catch (DecodingException e) {
-            LOGGER.info("Error during initial request parsing", e);
+            session.withMDC(() -> LOGGER.info("Error during initial request parsing", e));
             return unknownCommand(tag, session);
         }
     }

diff --git a/...rotocols-imap4/src/main/java/org/apache/james/imapserver/netty/HAProxyMessageHandler.java b/...rotocols-imap4/src/main/java/org/apache/james/imapserver/netty/HAProxyMessageHandler.java
@@ -19,7 +19,7 @@
 
 package org.apache.james.imapserver.netty;
 
-import static org.apache.james.imapserver.netty.ImapChannelUpstreamHandler.MDC_KEY;
+import static org.apache.james.imap.api.process.ImapSession.MDC_KEY;
 
 import java.net.InetSocketAddress;
 import java.util.Set;

diff --git a/...ocols/protocols-imap4/src/main/java/org/apache/james/imapserver/netty/IMAPMDCContext.java b/...ocols/protocols-imap4/src/main/java/org/apache/james/imapserver/netty/IMAPMDCContext.java
@@ -23,8 +23,6 @@
 import java.net.SocketAddress;
 import java.util.Optional;
 
-import org.apache.james.core.Username;
-import org.apache.james.imap.api.process.ImapSession;
 import org.apache.james.protocols.netty.ProtocolMDCContextFactory;
 import org.apache.james.util.MDCBuilder;
 
@@ -71,13 +69,4 @@ private static String retrieveHost(ChannelHandlerContext ctx) {
         }
         return remoteAddress.toString();
     }
-
-    public static MDCBuilder from(ImapSession imapSession) {
-        return MDCBuilder.create()
-            .addToContext(MDCBuilder.USER, Optional.ofNullable(imapSession.getUserName())
-                .map(Username::asString)
-                .orElse(""))
-            .addToContextIfPresent("selectedMailbox", Optional.ofNullable(imapSession.getSelected())
-                .map(selectedMailbox -> selectedMailbox.getMailboxId().serialize()));
-    }
 }
diff --git a/...ols-imap4/src/main/java/org/apache/james/imapserver/netty/ImapChannelUpstreamHandler.java b/...ols-imap4/src/main/java/org/apache/james/imapserver/netty/ImapChannelUpstreamHandler.java
@@ -18,6 +18,7 @@
  ****************************************************************/
 package org.apache.james.imapserver.netty;
 
+import static org.apache.james.imap.api.process.ImapSession.MDC_KEY;
 import static org.apache.james.imapserver.netty.IMAPServer.AuthenticationConfiguration;
 
 import java.io.Closeable;
@@ -78,7 +79,6 @@
 @ChannelHandler.Sharable
 public class ImapChannelUpstreamHandler extends ChannelInboundHandlerAdapter implements NettyConstants {
     private static final Logger LOGGER = LoggerFactory.getLogger(ImapChannelUpstreamHandler.class);
-    public static final String MDC_KEY = "bound_MDC";
 
     public static class ImapChannelUpstreamHandlerBuilder {
         private String hello;
@@ -254,12 +254,7 @@ private MDCBuilder mdc(ChannelHandlerContext ctx) {
 
     private MDCBuilder mdc(ImapSession imapSession) {
         return Optional.ofNullable(imapSession)
-            .map(session -> {
-                MDCBuilder boundMDC = (MDCBuilder) session.getAttribute(MDC_KEY);
-
-                return IMAPMDCContext.from(session)
-                    .addToContext(boundMDC);
-            })
+            .map(ImapSession::mdc)
             .orElseGet(MDCBuilder::create);
     }
 
@@ -297,7 +292,7 @@ static String retrieveUsername(ImapSession imapSession) {
     public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
         ImapSession imapSession = ctx.channel().attr(IMAP_SESSION_ATTRIBUTE_KEY).getAndSet(null);
         String username = retrieveUsername(imapSession);
-        try (Closeable closeable = mdc(ctx).build()) {
+        try (Closeable closeable = mdc(imapSession).build()) {
             if (cause instanceof SocketException) {
                 LOGGER.info("Socket exception encountered for user {}: {}", username, cause.getMessage());
             } else if (isSslHandshkeException(cause)) {

diff --git a/...ols/protocols-imap4/src/main/java/org/apache/james/imapserver/netty/NettyImapSession.java b/...ols/protocols-imap4/src/main/java/org/apache/james/imapserver/netty/NettyImapSession.java
@@ -30,6 +30,7 @@
 import javax.net.ssl.SSLSession;
 
 import org.apache.commons.lang3.NotImplementedException;
+import org.apache.james.core.Username;
 import org.apache.james.imap.api.ImapSessionState;
 import org.apache.james.imap.api.process.ImapLineHandler;
 import org.apache.james.imap.api.process.ImapSession;
@@ -40,6 +41,7 @@
 import org.apache.james.protocols.api.OidcSASLConfiguration;
 import org.apache.james.protocols.netty.Encryption;
 import org.apache.james.protocols.netty.LineHandlerAware;
+import org.apache.james.util.MDCBuilder;
 
 import io.netty.buffer.Unpooled;
 import io.netty.channel.Channel;
@@ -142,6 +144,7 @@ public Mono<Void> selected(SelectedMailbox mailbox) {
                 .orElse(Mono.empty()));
     }
 
+
     @Override
     public MailboxSession getMailboxSession() {
         return mailboxSession;
@@ -150,6 +153,25 @@ public MailboxSession getMailboxSession() {
     @Override
     public void setMailboxSession(MailboxSession mailboxSession) {
         this.mailboxSession = mailboxSession;
+        addUserToMDC(mailboxSession);
+    }
+
+    private void addUserToMDC(MailboxSession mailboxSession) {
+        setAttribute(MDC_KEY, mdc()
+            .addToContext(MDCBuilder.USER, Optional.ofNullable(mailboxSession.getUser())
+                .map(Username::asString)
+                .orElse("")));
+    }
+
+    @Override
+    public MDCBuilder mdc() {
+        SelectedMailbox mailbox = selectedMailbox.get();
+        if (mailbox != null) {
+            return MDCBuilder.create()
+                .addToContext(ImapSession.super.mdc())
+                .addToContext("selectedMailbox", mailbox.getMailboxId().serialize());
+        }
+        return ImapSession.super.mdc();
     }
 
     @Override