spb/zyquo-router Public MIT

1//2//  HTTPServer.swift3//  Zyquo Router4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//8//  Embedded HTTP/1.1 server on SwiftNIO's structured-concurrency APIs9//  (NIOAsyncChannel). One task per connection, one request at a time per10//  connection (keep-alive), spec-exact SSE via SSEWriter. Cancelling the11//  server task stops the accept loop and cancels in-flight connections,12//  which propagates into upstream provider calls.13//1415import Foundation16import NIOCore17import NIOFoundationCompat18import NIOHTTP119import NIOPosix2021/// What the server does with one parsed request.22/// `.complete` answers with a full buffered body; `.stream` hands an SSEWriter23/// to the route (Phase 3 streaming completions).24enum RouteResult {25    case complete(status: HTTPResponseStatus, headers: [(String, String)], body: Data)26    case stream(status: HTTPResponseStatus, headers: [(String, String)], body: (SSEWriter) async throws -> Void)27}2829/// One parsed inbound HTTP request, as handed to the route layer.30struct RouteRequest {31    let method: HTTPMethod32    let uri: String33    let headers: HTTPHeaders34    let body: Data3536    /// URI path without the query string.37    var path: String {38        uri.split(separator: "?", maxSplits: 1, omittingEmptySubsequences: false)39            .first.map(String.init) ?? uri40    }41}4243enum ServerError: LocalizedError {44    case portInUse(host: String, port: Int)45    case bindFailed(String)4647    var errorDescription: String? {48        switch self {49        case .portInUse(_, let port):50            return "Port \(port) is already in use — try \(port + 1), or stop the other process."51        case .bindFailed(let detail):52            return "Could not start the server: \(detail)"53        }54    }55}5657/// The embedded server. Create one per Start; it is single-use.58final class HTTPServer: Sendable {59    /// Maximum buffered request body (base64 images make chat bodies large).60    static let maxBodyBytes = 32 * 1024 * 10246162    let host: String63    let port: Int64    private let handler: @Sendable (RouteRequest) async -> RouteResult6566    init(host: String, port: Int, handler: @escaping @Sendable (RouteRequest) async -> RouteResult) {67        self.host = host68        self.port = port69        self.handler = handler70    }7172    /// Binds and serves until the surrounding task is cancelled.73    /// `onRunning` fires once the socket is bound and accepting.74    func run(onRunning: @escaping @Sendable () -> Void) async throws {75        let serverChannel: NIOAsyncChannel<NIOAsyncChannel<HTTPServerRequestPart, HTTPServerResponsePart>, Never>76        do {77            serverChannel = try await ServerBootstrap(group: MultiThreadedEventLoopGroup.singleton)78                .serverChannelOption(ChannelOptions.backlog, value: 64)79                .serverChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1)80                .childChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1)81                .bind(host: host, port: port) { channel in82                    channel.eventLoop.makeCompletedFuture {83                        try channel.pipeline.syncOperations.configureHTTPServerPipeline(withErrorHandling: true)84                        return try NIOAsyncChannel(wrappingChannelSynchronously: channel)85                    }86                }87        } catch let error as IOError where error.errnoCode == EADDRINUSE {88            throw ServerError.portInUse(host: host, port: port)89        } catch let error as NIOCore.ChannelError {90            throw ServerError.bindFailed(String(describing: error))91        }9293        onRunning()9495        // Cancellation of the surrounding task ends the accept iteration; the96        // task group then cancels and awaits every in-flight connection, so97        // shutdown is clean end-to-end. (withThrowingTaskGroup rather than a98        // discarding group: the package targets macOS 13.)99        try await serverChannel.executeThenClose { inbound, _ in100            try await withThrowingTaskGroup(of: Void.self) { group in101                for try await connection in inbound {102                    group.addTask { [handler] in103                        await Self.serve(connection: connection, handler: handler)104                    }105                }106            }107        }108    }109110    // MARK: - Per-connection loop111112    private static func serve(113        connection: NIOAsyncChannel<HTTPServerRequestPart, HTTPServerResponsePart>,114        handler: @Sendable (RouteRequest) async -> RouteResult115    ) async {116        try? await connection.executeThenClose { inbound, outbound in117            var iterator = inbound.makeAsyncIterator()118            while let part = try await iterator.next() {119                guard case .head(let head) = part else { continue }120121                var bodyBuffer = ByteBuffer()122                var tooLarge = false123                readLoop: while let next = try await iterator.next() {124                    switch next {125                    case .body(var chunk):126                        if bodyBuffer.readableBytes + chunk.readableBytes > maxBodyBytes {127                            tooLarge = true128                        } else {129                            bodyBuffer.writeBuffer(&chunk)130                        }131                    case .end:132                        break readLoop133                    case .head:134                        return  // protocol violation; drop the connection135                    }136                }137138                let keepAlive = head.isKeepAlive139                if tooLarge {140                    try await Self.write(141                        result: OpenAIError.response(142                            status: .payloadTooLarge,143                            message: "Request body exceeds the \(maxBodyBytes / (1024 * 1024)) MB limit.",144                            type: "invalid_request_error"145                        ),146                        version: head.version, keepAlive: false, outbound: outbound147                    )148                    return149                }150151                let request = RouteRequest(152                    method: head.method,153                    uri: head.uri,154                    headers: head.headers,155                    body: bodyBuffer.readData(length: bodyBuffer.readableBytes) ?? Data()156                )157                let result = await handler(request)158                try await Self.write(result: result, version: head.version, keepAlive: keepAlive, outbound: outbound)159                if !keepAlive { return }160            }161        }162    }163164    private static func write(165        result: RouteResult,166        version: HTTPVersion,167        keepAlive: Bool,168        outbound: NIOAsyncChannelOutboundWriter<HTTPServerResponsePart>169    ) async throws {170        switch result {171        case .complete(let status, let extraHeaders, let body):172            var headers = HTTPHeaders(extraHeaders)173            headers.replaceOrAdd(name: "Content-Length", value: String(body.count))174            headers.replaceOrAdd(name: "Connection", value: keepAlive ? "keep-alive" : "close")175            try await outbound.write(.head(HTTPResponseHead(version: version, status: status, headers: headers)))176            if !body.isEmpty {177                try await outbound.write(.body(.byteBuffer(ByteBuffer(bytes: body))))178            }179            try await outbound.write(.end(nil))180181        case .stream(let status, let extraHeaders, let body):182            var headers = HTTPHeaders(extraHeaders)183            headers.replaceOrAdd(name: "Content-Type", value: "text/event-stream")184            headers.replaceOrAdd(name: "Cache-Control", value: "no-cache")185            headers.replaceOrAdd(name: "Connection", value: "close")186            try await outbound.write(.head(HTTPResponseHead(version: version, status: status, headers: headers)))187            let writer = SSEWriter(outbound: outbound)188            try await body(writer)189            try await outbound.write(.end(nil))190        }191    }192}193