diff --git a/Package.swift b/Package.swift index d99564a..bf1fae8 100644 --- a/Package.swift +++ b/Package.swift @@ -68,7 +68,11 @@ let package = Package( name: "SwiftIMAPTests", dependencies: [ "SwiftMail", - .product(name: "Testing", package: "swift-testing") + .product(name: "Testing", package: "swift-testing"), + .product(name: "NIO", package: "swift-nio"), + .product(name: "NIOEmbedded", package: "swift-nio"), + .product(name: "NIOIMAP", package: "swift-nio-imap"), + .product(name: "Logging", package: "swift-log") ], resources: [ .copy("Resources") diff --git a/Sources/SwiftMail/IMAP/IMAP/Handler/XOAUTH2AuthenticationHandler.swift b/Sources/SwiftMail/IMAP/IMAP/Handler/XOAUTH2AuthenticationHandler.swift new file mode 100644 index 0000000..956a1da --- /dev/null +++ b/Sources/SwiftMail/IMAP/IMAP/Handler/XOAUTH2AuthenticationHandler.swift @@ -0,0 +1,117 @@ +import Foundation +import Logging +import NIO +import NIOIMAP +import NIOIMAPCore + +/// Handler responsible for managing the IMAP XOAUTH2 authentication exchange. +final class XOAUTH2AuthenticationHandler: BaseIMAPCommandHandler<[Capability]>, IMAPCommandHandler, @unchecked Sendable { + private var collectedCapabilities: [Capability] = [] + private var shouldSendCredentialsOnChallenge: Bool + private var credentials: ByteBuffer + private let serverLogger: Logger + private var lastServerError: String? + + init( + commandTag: String, + promise: EventLoopPromise<[Capability]>, + credentials: ByteBuffer, + expectsChallenge: Bool, + logger: Logger + ) { + self.credentials = credentials + self.shouldSendCredentialsOnChallenge = expectsChallenge + self.serverLogger = logger + super.init(commandTag: commandTag, promise: promise) + } + + override init(commandTag: String, promise: EventLoopPromise<[Capability]>) { + fatalError("Use init(commandTag:promise:credentials:expectsChallenge:logger:) instead") + } + + override func channelRead(context: ChannelHandlerContext, data: NIOAny) { + let response = unwrapInboundIn(data) + + if case .authenticationChallenge(var challengeBuffer) = response { + handleAuthenticationChallenge(&challengeBuffer, context: context) + } + + super.channelRead(context: context, data: data) + } + + private func handleAuthenticationChallenge(_ challenge: inout ByteBuffer, context: ChannelHandlerContext) { + let sendCredentials = lock.withLock { () -> Bool in + if shouldSendCredentialsOnChallenge { + shouldSendCredentialsOnChallenge = false + return true + } + return false + } + + if sendCredentials { + let credentialBuffer = credentials + credentials = context.channel.allocator.buffer(capacity: 0) + + context.channel + .writeAndFlush(IMAPClientHandler.OutboundIn.part(.continuationResponse(credentialBuffer))) + .cascadeFailure(to: promise) + return + } + + if let message = challenge.readString(length: challenge.readableBytes), !message.isEmpty { + lock.withLock { lastServerError = message } + serverLogger.error("XOAUTH2 server error: \(message)") + } else { + lock.withLock { lastServerError = nil } + } + + let emptyBuffer = context.channel.allocator.buffer(capacity: 0) + context.channel + .writeAndFlush(IMAPClientHandler.OutboundIn.part(.continuationResponse(emptyBuffer))) + .cascadeFailure(to: promise) + } + + override func handleTaggedOKResponse(_ response: TaggedResponse) { + super.handleTaggedOKResponse(response) + + let capabilities = lock.withLock { collectedCapabilities } + if !capabilities.isEmpty { + succeedWithResult(capabilities) + } else if case .ok(let responseText) = response.state, + let code = responseText.code, + case .capability(let caps) = code { + succeedWithResult(caps) + } else { + succeedWithResult([]) + } + } + + override func handleTaggedErrorResponse(_ response: TaggedResponse) { + let summary = String(describing: response.state) + let serverMessage = lock.withLock { lastServerError } + if let serverMessage, !serverMessage.isEmpty { + failWithError(IMAPError.authFailed("\(summary) (\(serverMessage))")) + } else { + failWithError(IMAPError.authFailed(summary)) + } + } + + override func handleError(_ error: Error) { + failWithError(error) + } + + override func handleUntaggedResponse(_ response: Response) -> Bool { + if super.handleUntaggedResponse(response) { + return true + } + + switch response { + case .untagged(.capabilityData(let capabilities)): + lock.withLock { collectedCapabilities = capabilities } + default: + break + } + + return false + } +} diff --git a/Sources/SwiftMail/IMAP/IMAPError.swift b/Sources/SwiftMail/IMAP/IMAPError.swift index 545d04e..5a5a4b1 100644 --- a/Sources/SwiftMail/IMAP/IMAPError.swift +++ b/Sources/SwiftMail/IMAP/IMAPError.swift @@ -20,6 +20,8 @@ public enum IMAPError: Error { case expungeFailed(String) case moveFailed(String) case commandNotSupported(String) + case authFailed(String) + case unsupportedAuthMechanism(String) } // Add CustomStringConvertible conformance for better error messages @@ -56,6 +58,10 @@ extension IMAPError: CustomStringConvertible { return "Move failed: \(reason)" case .commandNotSupported(let reason): return "Command not supported: \(reason)" + case .authFailed(let reason): + return "Authentication failed: \(reason)" + case .unsupportedAuthMechanism(let reason): + return "Unsupported authentication mechanism: \(reason)" } } } @@ -98,9 +104,13 @@ extension IMAPError: LocalizedError { return "Failed to move messages: \(reason)" case .commandNotSupported(let reason): return "The requested command is not supported by the server: \(reason)" + case .authFailed(let reason): + return "The IMAP authentication failed: \(reason)" + case .unsupportedAuthMechanism(let reason): + return "The server does not support the requested authentication mechanism: \(reason)" } } - + public var recoverySuggestion: String? { switch self { case .connectionFailed: @@ -117,8 +127,12 @@ extension IMAPError: LocalizedError { return "Make sure to select a mailbox before performing this operation." case .commandNotSupported: return "This operation may not be supported by your email provider." + case .authFailed: + return "Verify your OAuth credentials or request a fresh access token." + case .unsupportedAuthMechanism: + return "Check that your email provider supports XOAUTH2 for IMAP connections." default: return "Check the error details and try again." } } -} \ No newline at end of file +} \ No newline at end of file diff --git a/Sources/SwiftMail/IMAP/IMAPServer.swift b/Sources/SwiftMail/IMAP/IMAPServer.swift index dff7911..187ce2a 100644 --- a/Sources/SwiftMail/IMAP/IMAPServer.swift +++ b/Sources/SwiftMail/IMAP/IMAPServer.swift @@ -230,7 +230,7 @@ public actor IMAPServer { public func login(username: String, password: String) async throws { let command = LoginCommand(username: username, password: password) let loginCapabilities = try await executeCommand(command) - + // If we got capabilities from the login response, use them if !loginCapabilities.isEmpty { self.capabilities = Set(loginCapabilities) @@ -239,6 +239,89 @@ public actor IMAPServer { try await fetchCapabilities() } } + + /// Performs XOAUTH2 authentication for the current IMAP connection. + /// - Parameters: + /// - email: The full mailbox address to authenticate as. + /// - accessToken: The OAuth 2.0 access token. + /// - Throws: ``IMAPError.unsupportedAuthMechanism`` if the server does not advertise XOAUTH2 or ``IMAPError.authFailed`` when authentication fails. + public func authenticateXOAUTH2(email: String, accessToken: String) async throws { + let mechanism = AuthenticationMechanism("XOAUTH2") + let xoauthCapability = Capability.authenticate(mechanism) + + guard capabilities.contains(xoauthCapability) else { + throw IMAPError.unsupportedAuthMechanism("XOAUTH2 not advertised by server") + } + + // Ensure we have an active channel before proceeding + clearInvalidChannel() + + if self.channel == nil { + logger.info("Channel is nil, re-establishing connection before authentication") + try await connect() + } + + guard let channel = self.channel else { + throw IMAPError.connectionFailed("Channel not initialized") + } + + let expectsChallenge = !capabilities.contains(.saslIR) + let tag = generateCommandTag() + + let handlerPromise = channel.eventLoop.makePromise(of: [Capability].self) + var credentialBuffer = makeXOAUTH2InitialResponseBuffer(email: email, accessToken: accessToken) + let handler = XOAUTH2AuthenticationHandler( + commandTag: tag, + promise: handlerPromise, + credentials: credentialBuffer, + expectsChallenge: expectsChallenge, + logger: logger + ) + + try await channel.pipeline.addHandler(handler).get() + + let initialResponse: InitialResponse? + if expectsChallenge { + initialResponse = nil + } else { + credentialBuffer = makeXOAUTH2InitialResponseBuffer(email: email, accessToken: accessToken) + initialResponse = InitialResponse(credentialBuffer) + } + + let command = TaggedCommand(tag: tag, command: .authenticate(mechanism: mechanism, initialResponse: initialResponse)) + let wrapped = IMAPClientHandler.OutboundIn.part(CommandStreamPart.tagged(command)) + + let timeoutSeconds = 10 + let scheduledTask = group.next().scheduleTask(in: .seconds(Int64(timeoutSeconds))) { + self.logger.warning("XOAUTH2 authentication timed out after \(timeoutSeconds) seconds") + handlerPromise.fail(IMAPError.timeout) + } + + do { + try await channel.writeAndFlush(wrapped).get() + let capabilities = try await handlerPromise.futureResult.get() + + scheduledTask.cancel() + await handleConnectionTerminationInResponses(handler.untaggedResponses) + duplexLogger.flushInboundBuffer() + + if !capabilities.isEmpty { + self.capabilities = Set(capabilities) + } else { + try await fetchCapabilities() + } + } catch { + scheduledTask.cancel() + await handleConnectionTerminationInResponses(handler.untaggedResponses) + duplexLogger.flushInboundBuffer() + + if !handler.isCompleted { + try? await channel.pipeline.removeHandler(handler) + } + + throw error + } + } /// Identify the client to the server using the `ID` command. /// - Parameter identification: Information describing the client. Pass the default value to send no information. @@ -975,6 +1058,18 @@ public actor IMAPServer { } } } + + private func makeXOAUTH2InitialResponseBuffer(email: String, accessToken: String) -> ByteBuffer { + var buffer = ByteBufferAllocator().buffer(capacity: email.utf8.count + accessToken.utf8.count + 32) + buffer.writeString("user=") + buffer.writeString(email) + buffer.writeInteger(UInt8(0x01)) + buffer.writeString("auth=Bearer ") + buffer.writeString(accessToken) + buffer.writeInteger(UInt8(0x01)) + buffer.writeInteger(UInt8(0x01)) + return buffer + } /** Execute an IMAP command diff --git a/Tests/SwiftIMAPTests/XOAUTH2AuthenticationHandlerTests.swift b/Tests/SwiftIMAPTests/XOAUTH2AuthenticationHandlerTests.swift new file mode 100644 index 0000000..0bb5148 --- /dev/null +++ b/Tests/SwiftIMAPTests/XOAUTH2AuthenticationHandlerTests.swift @@ -0,0 +1,202 @@ +import Foundation +import Logging +import NIO +import NIOEmbedded +@preconcurrency import NIOIMAP +@preconcurrency import NIOIMAPCore +import Testing +@testable import SwiftMail + +struct XOAUTH2AuthenticationHandlerTests { + private let email = "user@example.com" + private let token = "ya29.A0AfH6SExample" + private let logger = Logger(label: "com.swiftmail.tests.xoauth2") + + @Test + func testSASLIRSuccess() async throws { + let (channel, promise, _) = try await setUpChannel(tag: "A001", expectsChallenge: false) + defer { _ = try? channel.finish() } + + let command = TaggedCommand( + tag: "A001", + command: .authenticate( + mechanism: AuthenticationMechanism("XOAUTH2"), + initialResponse: InitialResponse(makeCredentialBuffer(using: channel.allocator)) + ) + ) + + try await channel.writeAndFlush(IMAPClientHandler.OutboundIn.part(.tagged(command))) + + guard var outbound = try channel.readOutbound(as: ByteBuffer.self) else { + Issue.record("Expected AUTHENTICATE command") + return + } + let commandString = outbound.readString(length: outbound.readableBytes) + let expectedBase64 = makeBase64String() + #expect(commandString == "A001 AUTHENTICATE XOAUTH2 \(expectedBase64)\r\n") + + var okBuffer = channel.allocator.buffer(capacity: 0) + okBuffer.writeString("A001 OK AUTHENTICATE completed\r\n") + try channel.writeInbound(okBuffer) + + let capabilities = try await promise.futureResult.get() + #expect(capabilities.isEmpty) + } + + @Test + func testFallbackWithoutSASLIR() async throws { + let (channel, promise, _) = try await setUpChannel(tag: "A002", expectsChallenge: true) + defer { _ = try? channel.finish() } + + let command = TaggedCommand( + tag: "A002", + command: .authenticate( + mechanism: AuthenticationMechanism("XOAUTH2"), + initialResponse: nil + ) + ) + + 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) + #expect(firstLine == "A002 AUTHENTICATE XOAUTH2\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 data") + return + } + let continuationLine = continuation.readString(length: continuation.readableBytes) + let expectedBase64 = makeBase64String() + #expect(continuationLine == "\(expectedBase64)\r\n") + + var okBuffer = channel.allocator.buffer(capacity: 0) + okBuffer.writeString("A002 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) + defer { _ = try? channel.finish() } + + let command = TaggedCommand( + tag: "A003", + 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) // discard AUTH line + + var challengeBuffer = channel.allocator.buffer(capacity: 0) + challengeBuffer.writeString("+ eyJzdGF0dXMiOiI0MDEiLCJtZXNzYWdlIjoiSW52YWxpZCB0b2tlbiJ9\r\n") + try channel.writeInbound(challengeBuffer) + + guard var responseBuffer = try channel.readOutbound(as: ByteBuffer.self) else { + Issue.record("Expected empty continuation response") + return + } + let responseLine = responseBuffer.readString(length: responseBuffer.readableBytes) + #expect(responseLine == "\r\n") + + var noBuffer = channel.allocator.buffer(capacity: 0) + noBuffer.writeString("A003 NO AUTHENTICATE failed\r\n") + try channel.writeInbound(noBuffer) + + do { + _ = try await promise.futureResult.get() + Issue.record("Expected authentication failure") + } catch let error as IMAPError { + switch error { + case .authFailed(let message): + #expect(message.contains("AUTHENTICATE failed")) + default: + Issue.record("Unexpected IMAPError: \(error)") + } + } catch { + Issue.record("Unexpected error type: \(error)") + } + } + + @Test + func testDirectNOFailsAuthentication() async throws { + let (channel, promise, _) = try await setUpChannel(tag: "A004", expectsChallenge: false) + defer { _ = try? channel.finish() } + + let command = TaggedCommand( + tag: "A004", + 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) + + var noBuffer = channel.allocator.buffer(capacity: 0) + noBuffer.writeString("A004 NO AUTHENTICATE failed\r\n") + try channel.writeInbound(noBuffer) + + do { + _ = try await promise.futureResult.get() + Issue.record("Expected authentication failure") + } catch let error as IMAPError { + if case .authFailed = error { + // expected path + } 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()) + + let promise = channel.eventLoop.makePromise(of: [Capability].self) + let handler = XOAUTH2AuthenticationHandler( + commandTag: tag, + promise: promise, + credentials: makeCredentialBuffer(using: channel.allocator), + expectsChallenge: expectsChallenge, + logger: logger + ) + try await channel.pipeline.addHandler(handler) + + return (channel, promise, handler) + } + + private func makeCredentialBuffer(using allocator: ByteBufferAllocator) -> ByteBuffer { + var buffer = allocator.buffer(capacity: email.utf8.count + token.utf8.count + 32) + buffer.writeString("user=") + buffer.writeString(email) + buffer.writeInteger(UInt8(0x01)) + buffer.writeString("auth=Bearer ") + buffer.writeString(token) + buffer.writeInteger(UInt8(0x01)) + buffer.writeInteger(UInt8(0x01)) + return buffer + } + + private func makeBase64String() -> String { + let raw = "user=\(email)\u{01}auth=Bearer \(token)\u{01}\u{01}" + return Data(raw.utf8).base64EncodedString() + } +}