//===----------------------------------------------------------------------===//

//

// This source file is part of the SwiftNIO open source project

//

// Copyright (c) 2019-2024 Apple Inc. and the SwiftNIO project authors

// Licensed under Apache License v2.0

//

// See LICENSE.txt for license information

// See CONTRIBUTORS.txt for the list of SwiftNIO project authors

//

// SPDX-License-Identifier: Apache-2.0

//

//===----------------------------------------------------------------------===//

import NIOCore

/// Errors that may be raised by the `HTTPClientProtocolUpgrader`.

public struct NIOHTTPClientUpgradeError: Hashable, Error {

    // Uses the open enum style to allow additional errors to be added in future.

    private enum Code: Hashable {

        case responseProtocolNotFound

        case invalidHTTPOrdering

        case upgraderDeniedUpgrade

        case writingToHandlerDuringUpgrade

        case writingToHandlerAfterUpgradeCompleted

        case writingToHandlerAfterUpgradeFailed

        case receivedResponseBeforeRequestSent

        case receivedResponseAfterUpgradeCompleted

    }

    private var code: Code

    private init(_ code: Code) {

        self.code = code

    }

    public static let responseProtocolNotFound = NIOHTTPClientUpgradeError(.responseProtocolNotFound)

    public static let invalidHTTPOrdering = NIOHTTPClientUpgradeError(.invalidHTTPOrdering)

    public static let upgraderDeniedUpgrade = NIOHTTPClientUpgradeError(.upgraderDeniedUpgrade)

    public static let writingToHandlerDuringUpgrade = NIOHTTPClientUpgradeError(.writingToHandlerDuringUpgrade)

    public static let writingToHandlerAfterUpgradeCompleted = NIOHTTPClientUpgradeError(

        .writingToHandlerAfterUpgradeCompleted

    )

    public static let writingToHandlerAfterUpgradeFailed = NIOHTTPClientUpgradeError(

        .writingToHandlerAfterUpgradeFailed

    )

    public static let receivedResponseBeforeRequestSent = NIOHTTPClientUpgradeError(.receivedResponseBeforeRequestSent)

    public static let receivedResponseAfterUpgradeCompleted = NIOHTTPClientUpgradeError(

        .receivedResponseAfterUpgradeCompleted

    )

}

extension NIOHTTPClientUpgradeError: CustomStringConvertible {

    public var description: String {

        String(describing: self.code)

    }

}

/// An object that implements `NIOHTTPClientProtocolUpgrader` knows how to handle HTTP upgrade to

/// a protocol on a client-side channel.

/// It has the option of denying this upgrade based upon the server response.

public protocol NIOHTTPClientProtocolUpgrader {

    /// The protocol this upgrader knows how to support.

    var supportedProtocol: String { get }

    /// All the header fields the protocol requires in the request to successfully upgrade.

    /// These header fields will be added to the outbound request's "Connection" header field.

    /// It is the responsibility of the custom headers call to actually add these required headers.

    var requiredUpgradeHeaders: [String] { get }

    /// Additional headers to be added to the request, beyond the "Upgrade" and "Connection" headers.

    func addCustom(upgradeRequestHeaders: inout HTTPHeaders)

    /// Gives the receiving upgrader the chance to deny the upgrade based on the upgrade HTTP response.

    func shouldAllowUpgrade(upgradeResponse: HTTPResponseHead) -> Bool

    /// Called when the upgrade response has been flushed. At this time it is safe to mutate the channel

    /// pipeline to add whatever channel handlers are required.

    /// Until the returned `EventLoopFuture` succeeds, all received data will be buffered.

    func upgrade(context: ChannelHandlerContext, upgradeResponse: HTTPResponseHead) -> EventLoopFuture<Void>

}

/// A client-side channel handler that sends a HTTP upgrade handshake request to perform a HTTP-upgrade.

/// When the first HTTP request is sent, this handler will add all appropriate headers to perform an upgrade to

/// the a protocol. It may add headers for a set of protocols in preference order.

/// If the upgrade fails (i.e. response is not 101 Switching Protocols), this handler simply

/// removes itself from the pipeline. If the upgrade is successful, it upgrades the pipeline to the new protocol.

///

/// The request sends an order of preference to request which protocol it would like to use for the upgrade.

/// It will only upgrade to the protocol that is returned first in the list and does not currently

/// have the capability to upgrade to multiple simultaneous layered protocols.

public final class NIOHTTPClientUpgradeHandler: ChannelDuplexHandler, RemovableChannelHandler {

    public typealias OutboundIn = HTTPClientRequestPart

    public typealias OutboundOut = HTTPClientRequestPart

    public typealias InboundIn = HTTPClientResponsePart

    public typealias InboundOut = HTTPClientResponsePart

    private var upgraders: [NIOHTTPClientProtocolUpgrader]

    private let httpHandlers: [RemovableChannelHandler]

    private let upgradeCompletionHandler: (ChannelHandlerContext) -> Void

    /// Whether we've already seen the first response from our initial upgrade request.

    private var seenFirstResponse = false

    private var upgradeState: UpgradeState = .requestRequired

    private var receivedMessages: CircularBuffer<NIOAny> = CircularBuffer()

    /// Create a `HTTPClientUpgradeHandler`.

    ///

    /// - Parameter upgraders: All `HTTPClientProtocolUpgrader` objects that will add their upgrade request

    ///     headers and handle the upgrade if there is a response for their protocol. They should be placed in

    ///     order of the preference for the upgrade.

    /// - Parameter httpHandlers: All `RemovableChannelHandler` objects which will be removed from the pipeline

    ///     once the upgrade response is sent. This is used to ensure that the pipeline will be in a clean state

    ///     after the upgrade. It should include any handlers that are directly related to handling HTTP.

    ///     At the very least this should include the `HTTPEncoder` and `HTTPDecoder`, but should also include

    ///     any other handler that cannot tolerate receiving non-HTTP data.

    /// - Parameter upgradeCompletionHandler: A closure that will be fired when HTTP upgrade is complete.

    public convenience init(

        upgraders: [NIOHTTPClientProtocolUpgrader],

        httpHandlers: [RemovableChannelHandler],

        upgradeCompletionHandler: @escaping (ChannelHandlerContext) -> Void

    ) {

        self.init(_upgraders: upgraders, httpHandlers: httpHandlers, upgradeCompletionHandler: upgradeCompletionHandler)

    }

    private init(

        _upgraders upgraders: [NIOHTTPClientProtocolUpgrader],

        httpHandlers: [RemovableChannelHandler],

        upgradeCompletionHandler: @escaping (ChannelHandlerContext) -> Void

    ) {

        precondition(upgraders.count > 0, "A minimum of one protocol upgrader must be specified.")

        self.upgraders = upgraders

        self.httpHandlers = httpHandlers

        self.upgradeCompletionHandler = upgradeCompletionHandler

    }

    public func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {

        switch self.upgradeState {

        case .requestRequired:

            let updatedData = self.addHeadersToOutboundOut(data: data)

            context.write(updatedData, promise: promise)

        case .awaitingConfirmationResponse:

            // Still have full http stack.

            context.write(data, promise: promise)

        case .upgraderReady, .upgrading:

            promise?.fail(NIOHTTPClientUpgradeError.writingToHandlerDuringUpgrade)

            context.fireErrorCaught(NIOHTTPClientUpgradeError.writingToHandlerDuringUpgrade)

        case .upgradingAddingHandlers:

            // These are most likely messages immediately fired by a new protocol handler.

            // As that is added last we can just forward them on.

            context.write(data, promise: promise)

        case .upgradeComplete:

            // Upgrade complete and this handler should have been removed from the pipeline.

            promise?.fail(NIOHTTPClientUpgradeError.writingToHandlerAfterUpgradeCompleted)

            context.fireErrorCaught(NIOHTTPClientUpgradeError.writingToHandlerAfterUpgradeCompleted)

        case .upgradeFailed:

            // Upgrade failed and this handler should have been removed from the pipeline.

            promise?.fail(NIOHTTPClientUpgradeError.writingToHandlerAfterUpgradeCompleted)

            context.fireErrorCaught(NIOHTTPClientUpgradeError.writingToHandlerAfterUpgradeCompleted)

        }

    }

    private func addHeadersToOutboundOut(data: NIOAny) -> NIOAny {

        let interceptedOutgoingRequest = NIOHTTPClientUpgradeHandler.unwrapOutboundIn(data)

        if case .head(var requestHead) = interceptedOutgoingRequest {

            self.upgradeState = .awaitingConfirmationResponse

            self.addConnectionHeaders(to: &requestHead)

            self.addUpgradeHeaders(to: &requestHead)

            return NIOHTTPClientUpgradeHandler.wrapOutboundOut(.head(requestHead))

        }

        return data

    }

    private func addConnectionHeaders(to requestHead: inout HTTPRequestHead) {

        let requiredHeaders = ["upgrade"] + self.upgraders.flatMap { $0.requiredUpgradeHeaders }

        requestHead.headers.add(name: "Connection", value: requiredHeaders.joined(separator: ","))

    }

    private func addUpgradeHeaders(to requestHead: inout HTTPRequestHead) {

        let allProtocols = self.upgraders.map { $0.supportedProtocol.lowercased() }

        requestHead.headers.add(name: "Upgrade", value: allProtocols.joined(separator: ","))

        // Allow each upgrader the chance to add custom headers.

        for upgrader in self.upgraders {

            upgrader.addCustom(upgradeRequestHeaders: &requestHead.headers)

        }

    }

    public func channelRead(context: ChannelHandlerContext, data: NIOAny) {

        guard !self.seenFirstResponse else {

            // We're waiting for upgrade to complete: buffer this data.

            self.receivedMessages.append(data)

            return

        }

        let responsePart = NIOHTTPClientUpgradeHandler.unwrapInboundIn(data)

        switch self.upgradeState {

        case .awaitingConfirmationResponse:

            self.firstResponseHeadReceived(context: context, responsePart: responsePart)

        case .upgrading, .upgradingAddingHandlers:

            if case .end = responsePart {

                // This is the end of the first response. Swallow it, we're buffering the rest.

                self.seenFirstResponse = true

            }

        case .upgraderReady(let upgrade):

            if case .end = responsePart {

                // This is the end of the first response, and we can upgrade. Time to kick it off.

                self.seenFirstResponse = true

                upgrade()

            }

        case .upgradeFailed:

            // We were reentrantly called while delivering the response head. We can just pass this through.

            context.fireChannelRead(data)

        case .upgradeComplete:

            //Upgrade has completed but we have not seen a whole response and still got reentrantly called.

            context.fireErrorCaught(NIOHTTPClientUpgradeError.receivedResponseAfterUpgradeCompleted)

        case .requestRequired:

            //We are receiving an upgrade response and we have not requested the upgrade.

            context.fireErrorCaught(NIOHTTPClientUpgradeError.receivedResponseBeforeRequestSent)

        }

    }

    private func firstResponseHeadReceived(context: ChannelHandlerContext, responsePart: HTTPClientResponsePart) {

        // We should decide if we're can upgrade based on the first response header: if we aren't upgrading,

        // by the time the body comes in we should be out of the pipeline. That means that if we don't think we're

        // upgrading, the only thing we should see is a response head. Anything else in an error.

        guard case .head(let response) = responsePart else {

            self.notUpgrading(context: context, data: responsePart, error: .invalidHTTPOrdering)

            return

        }

        // Assess whether the upgrade response has accepted our upgrade request.

        guard case .switchingProtocols = response.status else {

            self.notUpgrading(context: context, data: responsePart, error: nil)

            return

        }

        do {

            let callback = try self.handleUpgrade(context: context, upgradeResponse: response)

            self.gotUpgrader(upgrader: callback)

        } catch {

            let clientError = error as? NIOHTTPClientUpgradeError

            self.notUpgrading(context: context, data: responsePart, error: clientError)

        }

    }

    private func handleUpgrade(

        context: ChannelHandlerContext,

        upgradeResponse response: HTTPResponseHead

    ) throws -> (() -> Void) {

        // Ok, we have a HTTP response. Check if it's an upgrade confirmation.

        // If it's not, we want to pass it on and remove ourselves from the channel pipeline.

        let acceptedProtocols = response.headers[canonicalForm: "upgrade"]

        // At the moment we only upgrade to the first protocol returned from the server.

        guard let protocolName = acceptedProtocols.first?.lowercased() else {

            // There are no upgrade protocols returned.

            throw NIOHTTPClientUpgradeError.responseProtocolNotFound

        }

        return try self.handleUpgradeForProtocol(

            context: context,

            protocolName: protocolName,

            response: response

        )

    }

    /// Attempt to upgrade a single protocol.

    private func handleUpgradeForProtocol(

        context: ChannelHandlerContext,

        protocolName: String,

        response: HTTPResponseHead

    ) throws -> (() -> Void) {

        let matchingUpgrader = self.upgraders

            .first(where: { $0.supportedProtocol.lowercased() == protocolName })

        guard let upgrader = matchingUpgrader else {

            // There is no upgrader for this protocol.

            throw NIOHTTPClientUpgradeError.responseProtocolNotFound

        }

        guard upgrader.shouldAllowUpgrade(upgradeResponse: response) else {

            // The upgrader says no.

            throw NIOHTTPClientUpgradeError.upgraderDeniedUpgrade

        }

        return self.performUpgrade(context: context, upgrader: upgrader, response: response)

    }

    private func performUpgrade(

        context: ChannelHandlerContext,

        upgrader: NIOHTTPClientProtocolUpgrader,

        response: HTTPResponseHead

    ) -> () -> Void {

        // Before we start the upgrade we have to remove the HTTPEncoder and HTTPDecoder handlers from the

        // pipeline, to prevent them parsing any more data. We'll buffer the incoming data until that completes.

        // While there are a lot of Futures involved here it's quite possible that all of this code will

        // actually complete synchronously: we just want to program for the possibility that it won't.

        // Once that's done, we call the internal handler, then call the upgrader code, and then finally when the

        // upgrader code is done, we do our final cleanup steps, namely we replay the received data we

        // buffered in the meantime and then remove ourselves from the pipeline.

        let pipeline = context.pipeline

        return {

            self.upgradeState = .upgrading

            self.removeHTTPHandlers(pipeline: pipeline)

                .map {

                    // Let the other handlers be removed before continuing with upgrade.

                    self.upgradeCompletionHandler(context)

                    self.upgradeState = .upgradingAddingHandlers

                }

                .flatMap {

                    upgrader.upgrade(context: context, upgradeResponse: response)

                }

                .map {

                    // We unbuffer any buffered data here.

                    // If we received any, we fire readComplete.

                    let fireReadComplete = self.receivedMessages.count > 0

                    while self.receivedMessages.count > 0 {

                        let bufferedPart = self.receivedMessages.removeFirst()

                        context.fireChannelRead(bufferedPart)

                    }

                    if fireReadComplete {

                        context.fireChannelReadComplete()

                    }

                    // We wait with the state change until _after_ the channel reads here.

                    // This is to prevent firing writes in response to these reads after we went to .upgradeComplete

                    // See: https://github.com/apple/swift-nio/issues/1279

                    self.upgradeState = .upgradeComplete

                }

                .whenComplete { _ in

                    context.pipeline.syncOperations.removeHandler(context: context, promise: nil)

                }

        }

    }

    /// Removes any extra HTTP-related handlers from the channel pipeline.

    private func removeHTTPHandlers(pipeline: ChannelPipeline) -> EventLoopFuture<Void>.Isolated {

        guard self.httpHandlers.count > 0 else {

            return pipeline.eventLoop.makeSucceededIsolatedFuture(())

        }

        let removeFutures = self.httpHandlers.map { pipeline.syncOperations.removeHandler($0) }

        return EventLoopFuture.andAllSucceed(removeFutures, on: pipeline.eventLoop).assumeIsolated()

    }

    private func gotUpgrader(upgrader: @escaping (() -> Void)) {

        self.upgradeState = .upgraderReady(upgrader)

        if self.seenFirstResponse {

            // Ok, we're good to go, we can upgrade. Otherwise we're waiting for .end, which

            // will trigger the upgrade.

            upgrader()

        }

    }

    private func notUpgrading(

        context: ChannelHandlerContext,

        data: HTTPClientResponsePart,

        error: NIOHTTPClientUpgradeError?

    ) {

        self.upgradeState = .upgradeFailed

        if let error = error {

            context.fireErrorCaught(error)

        }

        assert(self.receivedMessages.isEmpty)

        context.fireChannelRead(NIOHTTPClientUpgradeHandler.wrapInboundOut(data))

        // We've delivered the data. We can now remove ourselves, which should happen synchronously.

        context.pipeline.syncOperations.removeHandler(context: context, promise: nil)

    }

}

extension NIOHTTPClientUpgradeHandler: @unchecked Sendable {}

extension NIOHTTPClientUpgradeHandler {

    /// The state of the upgrade handler.

    fileprivate enum UpgradeState {

        /// Request not sent. This will need to be sent to initiate the upgrade.

        case requestRequired

        /// Awaiting confirmation response which will allow the upgrade to zero one or more protocols.

        case awaitingConfirmationResponse

        /// The response head has been received. We have an upgrader, which means we can begin upgrade.

        case upgraderReady(() -> Void)

        /// The response head has been received. The upgrade is in process.

        case upgrading

        /// The upgrade is in process and all of the http handlers have been removed.

        case upgradingAddingHandlers

        /// The upgrade has succeeded, and we are being removed from the pipeline.

        case upgradeComplete

        /// The upgrade has failed.

        case upgradeFailed

    }

}