Alternative version with a ContextTransformingMiddlewareProtocol

Author: Joannis

Sources/Hummingbird/Middleware/Middleware.swift

@@ -14,13 +14,18 @@
 
 import NIOCore
 
-/// Middleware protocol with generic input, context and output types
-public protocol MiddlewareProtocol<Input, Output, Context>: Sendable {
+public protocol ContextTransformingMiddlewareProtocol<Input, Output, InputContext, OutputContext>: Sendable {
     associatedtype Input
     associatedtype Output
-    associatedtype Context
+    associatedtype InputContext
+    associatedtype OutputContext
+
+    func handle(_ input: Input, context: InputContext, next: (Input, OutputContext) async throws -> Output) async throws -> Output
+}
 
-    func handle(_ input: Input, context: Context, next: (Input, Context) async throws -> Output) async throws -> Output
+/// Middleware protocol with generic input, context and output types
+public protocol MiddlewareProtocol<Input, Output, Context>: ContextTransformingMiddlewareProtocol where InputContext == OutputContext, Context == InputContext {
+    associatedtype Context
 }
 
 /// Applied to `Request` before it is dealt with by the router. Middleware passes the processed request onto the next responder

Sources/Hummingbird/Router/RouterMethods.swift

@@ -94,14 +94,14 @@ extension RouterMethods {
     /// ```
     /// - Parameters
     ///   - path: path prefix to add to routes inside this group
-    ///   - mapContext: Function converting context
+    ///   - middleware: A middleware that can transform the context
     @discardableResult public func group<TargetContext>(
         _ path: RouterPath = "",
-        mapContext: @escaping @Sendable (Request, Context) async throws -> TargetContext
+        middleware: some ContextTransformingMiddlewareProtocol<Request, Response, Context, TargetContext>
     ) -> RouterGroup<TargetContext> {
         return RouterGroup(
             path: path,
-            parent: TransformingRouterGroup(parent: self, transform: mapContext)
+            parent: TransformingRouterGroup(parent: self, middleware: middleware)
         )
     }
 

Sources/Hummingbird/Router/TransformingRouterGroup.swift

@@ -20,32 +20,38 @@ import NIOCore
 struct TransformingRouterGroup<InputContext: RequestContext, Context: RequestContext>: RouterMethods {
     typealias TransformContext = Context
     let parent: any RouterMethods<InputContext>
-    let transform: @Sendable (Request, InputContext) async throws -> TransformContext
+    let transform: @Sendable (Request, InputContext, (Request, TransformContext) async throws -> Response) async throws -> Response
 
     struct ContextTransformingResponder: HTTPResponder {
         typealias Context = InputContext
         let responder: any HTTPResponder<TransformContext>
-        let transform: @Sendable (Request, InputContext) async throws -> TransformContext
+        let transform: @Sendable (
+            Request,
+            InputContext,
+            (Request, TransformContext) async throws -> Response
+        ) async throws -> Response
 
         func respond(to request: Request, context: InputContext) async throws -> Response {
-            let newContext = try await transform(request, context)
-            return try await self.responder.respond(to: request, context: newContext)
+            try await transform(request, context) { req, context in
+                try await self.responder.respond(to: request, context: context)
+            }
         }
     }
 
     init(parent: any RouterMethods<InputContext>) where Context.Source == InputContext {
         self.parent = parent
-        self.transform = { _, context in
-            TransformContext(source: context)
+        self.transform = { req, context, next in
+            let context = TransformContext(source: context)
+            return try await next(req, context)
         }
     }
 
     init(
         parent: any RouterMethods<InputContext>,
-        transform: @escaping @Sendable (Request, InputContext) async throws -> TransformContext
+        middleware: some ContextTransformingMiddlewareProtocol<Request, Response, InputContext, Context>
     ) {
         self.parent = parent
-        self.transform = transform
+        self.transform = middleware.handle
     }
 
     /// Add middleware (Stub function as it isn't used)