Adding Timeout

Author: swhitty

Sources/HTTPConnection.swift

@@ -46,8 +46,8 @@ final class HTTPConnection: Hashable {
         HTTPRequestSequence(bytes: socket.bytes)
     }
 
-    func sendResponse(_ response: HTTPResponse, for request: HTTPRequest) throws {
-        try socket.write(HTTPResponseEncoder.encodeResponse(response, for: request))
+    func sendResponse(_ response: HTTPResponse) throws {
+        try socket.write(HTTPResponseEncoder.encodeResponse(response))
     }
 
     func close() throws {

Sources/HTTPResponse.swift

@@ -47,3 +47,9 @@ public struct HTTPResponse {
         self.body = body
     }
 }
+
+extension HTTPResponse {
+    var shouldKeepAlive: Bool {
+        headers[.connection]?.caseInsensitiveCompare("keep-alive") == .orderedSame
+    }
+}

Sources/HTTPResponseEncoder.swift

@@ -33,19 +33,15 @@ import Foundation
 
 struct HTTPResponseEncoder {
 
-    static func makeHeaderLines(from response: HTTPResponse, for request: HTTPRequest) -> [String] {
+    static func makeHeaderLines(from response: HTTPResponse) -> [String] {
         let status = [response.version.rawValue,
                       String(response.statusCode.code),
                       response.statusCode.phrase].joined(separator: " ")
 
         var httpHeaders = response.headers
 
-        if request.shouldKeepAlive {
-            httpHeaders[.connection] = "keep-alive"
-        }
-
         if response.body.isEmpty {
-            httpHeaders[.contentLength] = nil
+            httpHeaders[.contentLength] = String(0)
         } else {
             httpHeaders[.contentLength] = String(response.body.count)
         }
@@ -55,13 +51,15 @@ struct HTTPResponseEncoder {
         return [status] + headers + ["\r\n"]
     }
 
-    static func encodeResponse(_ response: HTTPResponse, for request: HTTPRequest) throws -> Data {
-        guard var data = makeHeaderLines(from: response, for: request)
+    static func encodeResponse(_ response: HTTPResponse) throws -> Data {
+        guard var data = makeHeaderLines(from: response)
                 .joined(separator: "\r\n")
                 .data(using: .utf8) else {
                     throw Error("Invalid Response Headers")
             }
+
         data.append(response.body)
+
         return data
     }
 }

Sources/HTTPServer.swift

@@ -34,12 +34,14 @@ import Foundation
 public final actor HTTPServer {
 
     private let port: UInt16
+    private let timeout: TimeInterval
     private var socket: AsyncSocket?
     private var connections = [HTTPConnection: Task<Void, Never>]()
     private var handlers: [(route: HTTPRoute, handler: HTTPHandler)]
 
-    public init(port: UInt16, handlers: [(route: HTTPRoute, handler: HTTPHandler)] = []) {
+    public init(port: UInt16, timeout: TimeInterval = 15, handlers: [(route: HTTPRoute, handler: HTTPHandler)] = []) {
         self.port = port
+        self.timeout = timeout
         self.handlers = []
     }
 
@@ -95,11 +97,12 @@ public final actor HTTPServer {
         connections[connection] = Task {
             do {
                 for try await request in connection.requests {
-                    let response = await self.handleRequest(request)
-                    try connection.sendResponse(response, for: request)
+                    let response = await handleRequest(request)
+                    try connection.sendResponse(response)
+                    guard response.shouldKeepAlive else { break }
                 }
             } catch {
-                print("error", connection.hostname, error)
+                print("connection error", connection.hostname, error)
             }
             removeConnection(connection)
         }
@@ -114,13 +117,24 @@ public final actor HTTPServer {
     }
 
     func handleRequest(_ request: HTTPRequest) async -> HTTPResponse {
+        var response = await handleRequest(request, timeout: timeout)
+        if request.shouldKeepAlive {
+            response.headers[.connection] = request.headers[.connection]
+        }
+        return response
+    }
+
+    func handleRequest(_ request: HTTPRequest, timeout: TimeInterval) async -> HTTPResponse {
         guard let handler = handlers.first(where: { $0.route ~= request })?.handler else {
             return HTTPResponse(statusCode: .notFound)
         }
 
         do {
-            return try await handler.handleRequest(request)
+            return try await withThrowingTimeout(seconds: timeout) {
+                try await handler.handleRequest(request)
+            }
         } catch {
+            print("handler error", error)
             return HTTPResponse(statusCode: .serverError)
         }
     }

Sources/Socket/Socket.swift

@@ -89,7 +89,6 @@ struct Socket: Hashable {
 
         if let address = listenAddress {
             guard address.withCString({ cstring in inet_pton(AF_INET6, cstring, &addr.sin6_addr) }) == 1 else {
-                print("\(address) is not converted.")
                 throw SocketError.bindFailed(makeErrorMessage())
             }
         }

Sources/Task+Timeout.swift

@@ -0,0 +1,61 @@
+//
+//  Task+Timeout.swift
+//  FlyingFox
+//
+//  Created by Simon Whitty on 15/02/2022.
+//  Copyright © 2022 Simon Whitty. All rights reserved.
+//
+//  Distributed under the permissive MIT license
+//  Get the latest version from here:
+//
+//  https://github.com/swhitty/FlyingFox
+//
+//  Permission is hereby granted, free of charge, to any person obtaining a copy
+//  of this software and associated documentation files (the "Software"), to deal
+//  in the Software without restriction, including without limitation the rights
+//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+//  copies of the Software, and to permit persons to whom the Software is
+//  furnished to do so, subject to the following conditions:
+//
+//  The above copyright notice and this permission notice shall be included in all
+//  copies or substantial portions of the Software.
+//
+//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+//  SOFTWARE.
+//
+
+import Foundation
+
+func withThrowingTimeout<T>(seconds: TimeInterval, body: @escaping @Sendable () async throws -> T) async throws -> T {
+    try await withThrowingTaskGroup(of: T.self) { group -> T in
+        group.addTask(operation: {
+            try await body()
+        })
+        group.addTask {
+            try await Task.sleep(nanoseconds: UInt64(seconds * 1_000_000_000))
+            throw TimeoutError()
+        }
+        let success = try await group.next()!
+        group.cancelAll()
+        return success
+    }
+}
+
+struct TimeoutError: LocalizedError {
+    var errorDescription: String? = "Timed out before completion"
+}
+
+extension Task where Failure == Error {
+
+    // Start a new Task with a timeout.
+    init(priority: TaskPriority? = nil, timeout: TimeInterval, operation: @escaping @Sendable () async throws -> Success) {
+        self = Task(priority: priority) {
+            try await withThrowingTimeout(seconds: timeout, body: operation)
+        }
+    }
+}

Tests/Task+TimeoutTests.swift

@@ -0,0 +1,80 @@
+//
+//  Task+TimeoutTests.swift
+//  FlyingFox
+//
+//  Created by Simon Whitty on 15/02/2022.
+//  Copyright © 2022 Simon Whitty. All rights reserved.
+//
+//  Distributed under the permissive MIT license
+//  Get the latest version from here:
+//
+//  https://github.com/swhitty/FlyingFox
+//
+//  Permission is hereby granted, free of charge, to any person obtaining a copy
+//  of this software and associated documentation files (the "Software"), to deal
+//  in the Software without restriction, including without limitation the rights
+//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+//  copies of the Software, and to permit persons to whom the Software is
+//  furnished to do so, subject to the following conditions:
+//
+//  The above copyright notice and this permission notice shall be included in all
+//  copies or substantial portions of the Software.
+//
+//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+//  SOFTWARE.
+//
+
+@testable import FlyingFox
+import XCTest
+
+final class TaskTimeoutTests: XCTestCase {
+
+    func testTimeoutReturnsSuccess_WhenTimeoutDoesNotExpire() async throws {
+        // given
+        let value = try await Task(timeout: 0.5) {
+            "Fish"
+        }.value
+
+        // then
+        XCTAssertEqual(value, "Fish")
+    }
+
+    func testTimeoutThrowsError_WhenTimeoutExpires() async {
+        // given
+        let task = Task<Void, Error>(timeout: 0.5) {
+            try? await Task.sleep(nanoseconds: 10_000_000_000)
+        }
+
+        // then
+        do {
+            _ = try await task.value
+            XCTFail("Expected TimeoutError")
+        } catch {
+            XCTAssertTrue(error is TimeoutError)
+        }
+    }
+
+
+    func testTimeoutCancels() async {
+        // given
+        let task = Task(timeout: 0.5) {
+            try? await Task.sleep(nanoseconds: 10_000_000_000)
+        }
+
+        // when
+        task.cancel()
+
+        // then
+        do {
+            _ = try await task.value
+            XCTFail("Expected CancellationError")
+        } catch {
+            XCTAssertTrue(error is CancellationError)
+        }
+    }
+}