From b3efc762e9ef2634ab01a4fcb35566e095968ecb Mon Sep 17 00:00:00 2001
From: Adam Fowler <adamfowler71@gmail.com>
Date: Wed, 22 May 2024 17:26:22 +0100
Subject: [PATCH] Ensure errors have server and date header

---
 Sources/Hummingbird/Application.swift         | 17 +++++++++++++-
 .../Error/NIOCore+HTTPResponseError.swift     | 23 +++++++++++++++++++
 .../Server/HTTP/HTTPChannelHandler.swift      |  7 ------
 Tests/HummingbirdTests/ApplicationTests.swift | 14 +++++++++++
 4 files changed, 53 insertions(+), 8 deletions(-)
 create mode 100644 Sources/HummingbirdCore/Error/NIOCore+HTTPResponseError.swift

diff --git a/Sources/Hummingbird/Application.swift b/Sources/Hummingbird/Application.swift
index 7d69dfec0..338aac50c 100644
--- a/Sources/Hummingbird/Application.swift
+++ b/Sources/Hummingbird/Application.swift
@@ -105,7 +105,22 @@ extension ApplicationProtocol {
                 logger: self.logger.with(metadataKey: "hb_id", value: .stringConvertible(RequestID()))
             )
             // respond to request
-            var response = try await responder.respond(to: request, context: context)
+            var response: Response
+            do {
+                response = try await responder.respond(to: request, context: context)
+            } catch {
+                switch error {
+                case let httpError as HTTPResponseError:
+                    // this is a processed error so don't log as Error
+                    response = httpError.response(allocator: channel.allocator)
+                default:
+                    // this error has not been recognised
+                    response = Response(
+                        status: .internalServerError,
+                        body: .init()
+                    )
+                }
+            }
             response.headers[.date] = dateCache.date
             // server name header
             if let serverName = self.configuration.serverName {
diff --git a/Sources/HummingbirdCore/Error/NIOCore+HTTPResponseError.swift b/Sources/HummingbirdCore/Error/NIOCore+HTTPResponseError.swift
new file mode 100644
index 000000000..22449b00f
--- /dev/null
+++ b/Sources/HummingbirdCore/Error/NIOCore+HTTPResponseError.swift
@@ -0,0 +1,23 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Hummingbird server framework project
+//
+// Copyright (c) 2024 the Hummingbird authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See hummingbird/CONTRIBUTORS.txt for the list of Hummingbird authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import HTTPTypes
+import NIOCore
+
+// If we catch a too many bytes error report that as payload too large
+extension NIOTooManyBytesError: HTTPResponseError {
+    public var status: HTTPResponse.Status { .contentTooLarge }
+    public var headers: HTTPFields { [:] }
+    public func body(allocator: ByteBufferAllocator) -> ByteBuffer? { nil }
+}
diff --git a/Sources/HummingbirdCore/Server/HTTP/HTTPChannelHandler.swift b/Sources/HummingbirdCore/Server/HTTP/HTTPChannelHandler.swift
index 14e97715a..13838e556 100644
--- a/Sources/HummingbirdCore/Server/HTTP/HTTPChannelHandler.swift
+++ b/Sources/HummingbirdCore/Server/HTTP/HTTPChannelHandler.swift
@@ -145,13 +145,6 @@ struct HTTPServerBodyWriter: Sendable, ResponseBodyWriter {
     }
 }
 
-// If we catch a too many bytes error report that as payload too large
-extension NIOTooManyBytesError: HTTPResponseError {
-    public var status: HTTPResponse.Status { .contentTooLarge }
-    public var headers: HTTPFields { [:] }
-    public func body(allocator: ByteBufferAllocator) -> ByteBuffer? { nil }
-}
-
 extension NIOLockedValueBox {
     /// Exchange stored value for new value and return the old stored value
     func exchange(_ newValue: Value) -> Value {
diff --git a/Tests/HummingbirdTests/ApplicationTests.swift b/Tests/HummingbirdTests/ApplicationTests.swift
index c1fc2063e..922e65f03 100644
--- a/Tests/HummingbirdTests/ApplicationTests.swift
+++ b/Tests/HummingbirdTests/ApplicationTests.swift
@@ -208,6 +208,20 @@ final class ApplicationTests: XCTestCase {
         }
     }
 
+    func testErrorHeaders() async throws {
+        let router = Router()
+        router.get("error") { _, _ -> HTTPResponse.Status in
+            throw HTTPError(.badRequest, message: "BAD!")
+        }
+        let app = Application(router: router, configuration: .init(serverName: "HB"))
+        try await app.test(.live) { client in
+            try await client.execute(uri: "/error", method: .get) { response in
+                XCTAssertEqual(response.headers[.server], "HB")
+                XCTAssertNotNil(response.headers[.date])
+            }
+        }
+    }
+
     func testResponseBody() async throws {
         let router = Router()
         router