Merge pull request #90 from jdudley/codex/pr-microsoft-xoauth2-upstream

IMAP XOAUTH2 interoperability hardening for Microsoft 365

Sources/SwiftMail/IMAP/IMAP/Handler/CommandHandler.swift

@@ -70,12 +70,28 @@ class BaseIMAPCommandHandler<ResultType: Sendable>: CommandHandler, RemovableCha
     /// Succeed the promise with a result
     /// - Parameter result: The result to succeed with
     func succeedWithResult(_ result: ResultType) {
+        let shouldSucceed = lock.withLock { () -> Bool in
+            if isCompleted {
+                return false
+            }
+            isCompleted = true
+            return true
+        }
+        guard shouldSucceed else { return }
         promise.succeed(result)
     }
     
     /// Fail the promise with an error
     /// - Parameter error: The error to fail with
     func failWithError(_ error: Error) {
+        let shouldFail = lock.withLock { () -> Bool in
+            if isCompleted {
+                return false
+            }
+            isCompleted = true
+            return true
+        }
+        guard shouldFail else { return }
         promise.fail(error)
     }
     
@@ -181,7 +197,17 @@ class BaseIMAPCommandHandler<ResultType: Sendable>: CommandHandler, RemovableCha
         // Always forward the response to the next handler
         context.fireChannelRead(data)
     }
-    
+    /// Channel inactive method from ChannelInboundHandler
+    func channelInactive(context: ChannelHandlerContext) {
+        let shouldFail = lock.withLock { !isCompleted }
+        if shouldFail {
+            failWithError(IMAPError.connectionFailed("Connection closed before command completed"))
+            handleCompletion(context: context)
+        }
+
+        context.fireChannelInactive()
+    }
+
     /// Error caught method from ChannelInboundHandler
     func errorCaught(context: ChannelHandlerContext, error: Error) {
         // Handle the error

Sources/SwiftMail/IMAP/IMAP/Handler/XOAUTH2AuthenticationHandler.swift

@@ -9,8 +9,10 @@ final class XOAUTH2AuthenticationHandler: BaseIMAPCommandHandler<[Capability]>,
     private var collectedCapabilities: [Capability] = []
     private var shouldSendCredentialsOnChallenge: Bool
     private var credentials: ByteBuffer
+    private let sentInlineInitialResponse: Bool
     private let serverLogger: Logger
     private var lastServerError: String?
+    private var fallbackContinuationSent = false
 
     init(
         commandTag: String,
@@ -22,6 +24,7 @@ final class XOAUTH2AuthenticationHandler: BaseIMAPCommandHandler<[Capability]>,
         self.credentials = credentials
         self.shouldSendCredentialsOnChallenge = expectsChallenge
         self.serverLogger = logger
+        self.sentInlineInitialResponse = !expectsChallenge
         super.init(commandTag: commandTag, promise: promise)
     }
 
@@ -40,11 +43,21 @@ final class XOAUTH2AuthenticationHandler: BaseIMAPCommandHandler<[Capability]>,
     }
 
     private func handleAuthenticationChallenge(_ challenge: inout ByteBuffer, context: ChannelHandlerContext) {
+        let challengeText = challenge.getString(at: challenge.readerIndex, length: challenge.readableBytes) ?? ""
+        let challengeIsEmpty = challengeText.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
+
         let sendCredentials = lock.withLock { () -> Bool in
             if shouldSendCredentialsOnChallenge {
                 shouldSendCredentialsOnChallenge = false
                 return true
             }
+
+            // Compatibility fallback: some servers advertise SASL-IR but still emit an
+            // empty continuation before consuming credentials. Allow one retry.
+            if sentInlineInitialResponse && !fallbackContinuationSent && challengeIsEmpty {
+                fallbackContinuationSent = true
+                return true
+            }
             return false
         }
 
@@ -58,9 +71,9 @@ final class XOAUTH2AuthenticationHandler: BaseIMAPCommandHandler<[Capability]>,
             return
         }
 
-        if let message = challenge.readString(length: challenge.readableBytes), !message.isEmpty {
-            lock.withLock { lastServerError = message }
-            serverLogger.error("XOAUTH2 server error: \(message)")
+        if !challengeText.isEmpty {
+            lock.withLock { lastServerError = challengeText }
+            serverLogger.error("XOAUTH2 server error: \(challengeText)")
         } else {
             lock.withLock { lastServerError = nil }
         }

Sources/SwiftMail/IMAP/IMAPConnection.swift

@@ -424,11 +424,19 @@ final class IMAPConnection {
             throw IMAPError.connectionFailed("Channel not initialized")
         }
 
-        let expectsChallenge = !capabilities.contains(.saslIR)
         let tag = generateCommandTag()
 
         let handlerPromise = channel.eventLoop.makePromise(of: [Capability].self)
         let credentialBuffer = makeXOAUTH2InitialResponseBuffer(email: email, accessToken: accessToken)
+
+        let supportsSASLIR = capabilities.contains(.saslIR)
+        let saslIRInlineLimitBytes = 1024
+        let shouldUseInlineInitialResponse = supportsSASLIR && credentialBuffer.readableBytes <= saslIRInlineLimitBytes
+        let expectsChallenge = !shouldUseInlineInitialResponse
+
+        if supportsSASLIR && !shouldUseInlineInitialResponse {
+            logger.info("XOAUTH2 payload size \(credentialBuffer.readableBytes) exceeds inline SASL-IR limit \(saslIRInlineLimitBytes); switching to continuation mode")
+        }
         let handler = XOAUTH2AuthenticationHandler(
             commandTag: tag,
             promise: handlerPromise,
@@ -440,14 +448,24 @@ final class IMAPConnection {
         try await channel.pipeline.addHandler(handler, position: .before(responseBuffer)).get()
         responseBuffer.hasActiveHandler = true
 
-        let initialResponse = expectsChallenge ? nil : InitialResponse(credentialBuffer)
+        let initialResponse: InitialResponse?
+        if shouldUseInlineInitialResponse {
+            initialResponse = InitialResponse(credentialBuffer)
+        } else if supportsSASLIR {
+            // Use true continuation mode for oversized payloads.
+            // Sending an explicit empty SASL-IR ("=") can be interpreted as an empty credential attempt.
+            initialResponse = nil
+        } else {
+            initialResponse = nil
+        }
 
         let command = TaggedCommand(tag: tag, command: .authenticate(mechanism: mechanism, initialResponse: initialResponse))
         let wrapped = IMAPClientHandler.OutboundIn.part(CommandStreamPart.tagged(command))
 
         let authenticationTimeoutSeconds = 10
         let logger = self.logger
-        let scheduledTask = group.next().scheduleTask(in: .seconds(Int64(authenticationTimeoutSeconds))) {
+        // Schedule on the channel event loop to avoid cross-loop promise completion.
+        let scheduledTask = channel.eventLoop.scheduleTask(in: .seconds(Int64(authenticationTimeoutSeconds))) {
             logger.warning("XOAUTH2 authentication timed out after \(authenticationTimeoutSeconds) seconds")
             handlerPromise.fail(IMAPError.timeout)
         }
@@ -462,7 +480,14 @@ final class IMAPConnection {
             duplexLogger.flushInboundBuffer()
 
             isSessionAuthenticated = true
-            try await refreshCapabilities(using: refreshedCapabilities)
+            if !refreshedCapabilities.isEmpty {
+                self.capabilities = Set(refreshedCapabilities)
+            } else {
+                // AUTHENTICATE often returns an OK without CAPABILITY data.
+                // Avoid issuing a follow-up CAPABILITY command here because we're already
+                // inside commandQueue.run, and a nested executeCommand would deadlock.
+                logger.debug("XOAUTH2 completed without capability data; retaining existing capability snapshot")
+            }
         } catch {
             scheduledTask.cancel()
             responseBuffer.hasActiveHandler = false
@@ -689,7 +714,7 @@ final class IMAPConnection {
         let timeoutSeconds = command.timeoutSeconds
 
         let logger = self.logger
-        let scheduledTask = group.next().scheduleTask(in: .seconds(Int64(timeoutSeconds))) {
+        let scheduledTask = channel.eventLoop.scheduleTask(in: .seconds(Int64(timeoutSeconds))) {
             logger.warning("Command timed out after \(timeoutSeconds) seconds")
             resultPromise.fail(IMAPError.timeout)
         }
@@ -743,7 +768,7 @@ final class IMAPConnection {
         let handler = HandlerType.init(commandTag: "", promise: resultPromise)
 
         let logger = self.logger
-        let scheduledTask = group.next().scheduleTask(in: .seconds(Int64(timeoutSeconds))) {
+        let scheduledTask = channel.eventLoop.scheduleTask(in: .seconds(Int64(timeoutSeconds))) {
             logger.warning("Handler execution timed out after \(timeoutSeconds) seconds")
             resultPromise.fail(IMAPError.timeout)
         }

Tests/SwiftIMAPTests/XOAUTH2AuthenticationHandlerTests.swift

@@ -85,6 +85,49 @@ struct XOAUTH2AuthenticationHandlerTests {
         #expect(capabilities.isEmpty)
     }
 
+
+    @Test
+    func testSASLIRServerSendsEmptyChallengeRetriesCredentials() async throws {
+        let (channel, promise, _) = try await setUpChannel(tag: "A002A", expectsChallenge: false)
+        defer { _ = try? channel.finish() }
+
+        let command = TaggedCommand(
+            tag: "A002A",
+            command: .authenticate(
+                mechanism: AuthenticationMechanism("XOAUTH2"),
+                initialResponse: InitialResponse(makeCredentialBuffer(using: channel.allocator))
+            )
+        )
+
+        try await channel.writeAndFlush(IMAPClientHandler.OutboundIn.part(.tagged(command)))
+
+        guard var firstOutbound = try channel.readOutbound(as: ByteBuffer.self) else {
+            Issue.record("Expected AUTHENTICATE command")
+            return
+        }
+        let firstLine = firstOutbound.readString(length: firstOutbound.readableBytes)
+        let expectedBase64 = makeBase64String()
+        #expect(firstLine == "A002A AUTHENTICATE XOAUTH2 \(expectedBase64)\r\n")
+
+        var challengeBuffer = channel.allocator.buffer(capacity: 0)
+        challengeBuffer.writeString("+ \r\n")
+        try channel.writeInbound(challengeBuffer)
+
+        guard var continuation = try channel.readOutbound(as: ByteBuffer.self) else {
+            Issue.record("Expected XOAUTH2 continuation retry data")
+            return
+        }
+        let continuationLine = continuation.readString(length: continuation.readableBytes)
+        #expect(continuationLine == "\(expectedBase64)\r\n")
+
+        var okBuffer = channel.allocator.buffer(capacity: 0)
+        okBuffer.writeString("A002A OK AUTHENTICATE completed\r\n")
+        try channel.writeInbound(okBuffer)
+
+        let capabilities = try await promise.futureResult.get()
+        #expect(capabilities.isEmpty)
+    }
+
     @Test
     func testServerErrorBlobTriggersAuthFailure() async throws {
         let (channel, promise, _) = try await setUpChannel(tag: "A003", expectsChallenge: false)
@@ -166,6 +209,38 @@ struct XOAUTH2AuthenticationHandlerTests {
         }
     }
 
+
+    @Test
+    func testChannelCloseFailsPendingAuthentication() async throws {
+        let (channel, promise, _) = try await setUpChannel(tag: "A005", expectsChallenge: false)
+
+        let command = TaggedCommand(
+            tag: "A005",
+            command: .authenticate(
+                mechanism: AuthenticationMechanism("XOAUTH2"),
+                initialResponse: InitialResponse(makeCredentialBuffer(using: channel.allocator))
+            )
+        )
+
+        try await channel.writeAndFlush(IMAPClientHandler.OutboundIn.part(.tagged(command)))
+        _ = try channel.readOutbound(as: ByteBuffer.self)
+
+        try await channel.close().get()
+
+        do {
+            _ = try await promise.futureResult.get()
+            Issue.record("Expected connection failure when channel closes")
+        } catch let error as IMAPError {
+            if case .connectionFailed(let message) = error {
+                #expect(message.contains("Connection closed before command completed"))
+            } else {
+                Issue.record("Unexpected IMAPError: \(error)")
+            }
+        } catch {
+            Issue.record("Unexpected error type: \(error)")
+        }
+    }
+
     private func setUpChannel(tag: String, expectsChallenge: Bool) async throws -> (EmbeddedChannel, EventLoopPromise<[Capability]>, XOAUTH2AuthenticationHandler) {
         let channel = EmbeddedChannel()
         try await channel.pipeline.addHandler(IMAPClientHandler())