ccn/Packages/Persistence/Tests/Utils/Fakes/PlaintextPostgresServer.swift

import NIOCore
import NIOPosix

/// A fake PostgreSQL server speaking just enough of the wire protocol to complete a plaintext startup.
///
/// It answers the client's `SSLRequest` with `'N'` (no SSL) and the subsequent startup message with
/// `AuthenticationOk`, `BackendKeyData`, and `ReadyForQuery` — so a client asking for TLS can only end up
/// connected in plaintext. This is what the `prefer` fallback test connects to, proving the driver
/// downgrades to plaintext rather than refusing the connection, and what the `require` test connects to,
/// proving the driver refuses the connection instead of downgrading.
final class PlaintextPostgresServer {

    // MARK: Properties

    /// The port the server listens on, assigned by the system at bind time.
    let port: Int

    /// The listening channel the server accepts connections through.
    private let channel: Channel

    // MARK: Initializers

    private init(
        channel: Channel,
        port: Int
    ) {
        self.channel = channel
        self.port = port
    }

    // MARK: Functions

    /// Starts a server on the loopback interface, on a system-assigned port.
    /// - Returns: the running server, ready to be connected to at ``port``.
    static func start() async throws -> PlaintextPostgresServer {
        let channel = try await ServerBootstrap(group: MultiThreadedEventLoopGroup.singleton)
            .childChannelInitializer { channel in
                channel.eventLoop.makeCompletedFuture {
                    try channel.pipeline.syncOperations.addHandler(Handler())
                }
            }
            .bind(host: "127.0.0.1", port: 0)
            .get()

        guard let port = channel.localAddress?.port else {
            throw ChannelError.unknownLocalAddress
        }

        return .init(
            channel: channel,
            port: port
        )
    }

    /// Stops the server, closing its listening channel.
    func stop() async throws {
        try await channel.close().get()
    }

}

// MARK: - Handlers

private extension PlaintextPostgresServer {

    /// Refuses the client's `SSLRequest`, accepts whatever startup message arrives, and closes on anything
    /// after that (e.g. a `Terminate`).
    ///
    /// Unlike MySQL, the PostgreSQL client speaks first, so nothing is written on `channelActive`.
    final class Handler: ChannelInboundHandler {

        // MARK: Type aliases

        typealias InboundIn = ByteBuffer
        typealias OutboundOut = ByteBuffer

        // MARK: Enumerations

        /// The startup phases the connection moves through.
        private enum State {
            case awaitingSSLRequest
            case awaitingStartup
            case established
        }

        // MARK: Constants

        /// The magic code identifying an `SSLRequest` message.
        private static let sslRequestCode: Int32 = 80877103

        // MARK: Properties

        /// The startup phase the connection is currently in.
        private var state: State = .awaitingSSLRequest

        // MARK: Functions

        func channelRead(
            context: ChannelHandlerContext,
            data: NIOAny
        ) {
            let buffer = unwrapInboundIn(data)

            switch state {
            case .awaitingSSLRequest:
                // Peek past the Int32 length at the Int32 code: an `SSLRequest` is refused with a bare
                // 'N', while a direct startup message (a client connecting with TLS disabled) is
                // answered straight away.
                guard buffer.getInteger(at: buffer.readerIndex + 4, as: Int32.self) == Self.sslRequestCode else {
                    completeStartup(context: context)
                    return
                }

                state = .awaitingStartup

                var refusal = context.channel.allocator.buffer(capacity: 1)

                refusal.writeInteger(UInt8(ascii: "N"))

                context.writeAndFlush(
                    wrapOutboundOut(refusal),
                    promise: nil
                )
            case .awaitingStartup:
                completeStartup(context: context)
            case .established:
                context.close(promise: nil)
            }
        }

        /// Answers a startup message and marks the connection established.
        private func completeStartup(context: ChannelHandlerContext) {
            state = .established

            context.writeAndFlush(
                wrapOutboundOut(Self.startupResponse(allocator: context.channel.allocator)),
                promise: nil
            )
        }

        // MARK: Helpers

        /// The reply completing a plaintext startup: `AuthenticationOk`, `BackendKeyData`, and
        /// `ReadyForQuery` in a single flush.
        ///
        /// `BackendKeyData` is not optional filler — the client requires it before `ReadyForQuery` by
        /// default and fails the connection when it is missing.
        private static func startupResponse(allocator: ByteBufferAllocator) -> ByteBuffer {
            var buffer = allocator.buffer(capacity: 32)

            buffer.writeInteger(UInt8(ascii: "R"))   // AuthenticationOk
            buffer.writeInteger(Int32(8))
            buffer.writeInteger(Int32(0))

            buffer.writeInteger(UInt8(ascii: "K"))   // BackendKeyData
            buffer.writeInteger(Int32(12))
            buffer.writeInteger(Int32(1))            // process id
            buffer.writeInteger(Int32(0))            // secret key

            buffer.writeInteger(UInt8(ascii: "Z"))   // ReadyForQuery
            buffer.writeInteger(Int32(5))
            buffer.writeInteger(UInt8(ascii: "I"))   // idle

            return buffer
        }

    }

}