Add IMAP XOAUTH2 authentication handler and tests

2026-03-17 20:02:25 +00:00 · 2025-09-18 11:14:35 +02:00
parent bd0be2c8a7
commit 544efd8c5a
5 changed files with 436 additions and 4 deletions
@@ -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")
@@ -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
+    }
+}
@@ -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."
        }
    }
-} 
+}
@@ -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
@@ -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()
+    }
+}