diff --git a/Sources/SwiftMail/IMAP/IMAP/Handler/CommandHandler.swift b/Sources/SwiftMail/IMAP/IMAP/Handler/CommandHandler.swift index 4277932..07837a2 100644 --- a/Sources/SwiftMail/IMAP/IMAP/Handler/CommandHandler.swift +++ b/Sources/SwiftMail/IMAP/IMAP/Handler/CommandHandler.swift @@ -70,12 +70,28 @@ class BaseIMAPCommandHandler: 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: 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 diff --git a/Sources/SwiftMail/IMAP/IMAP/Handler/XOAUTH2AuthenticationHandler.swift b/Sources/SwiftMail/IMAP/IMAP/Handler/XOAUTH2AuthenticationHandler.swift index 956a1da..2e46f1e 100644 --- a/Sources/SwiftMail/IMAP/IMAP/Handler/XOAUTH2AuthenticationHandler.swift +++ b/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 } } diff --git a/Sources/SwiftMail/IMAP/IMAPConnection.swift b/Sources/SwiftMail/IMAP/IMAPConnection.swift index f8bf0a0..c96219c 100644 --- a/Sources/SwiftMail/IMAP/IMAPConnection.swift +++ b/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) } diff --git a/Tests/SwiftIMAPTests/XOAUTH2AuthenticationHandlerTests.swift b/Tests/SwiftIMAPTests/XOAUTH2AuthenticationHandlerTests.swift index 0bb5148..fad77ed 100644 --- a/Tests/SwiftIMAPTests/XOAUTH2AuthenticationHandlerTests.swift +++ b/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())