Make all channel handlers removeable (#8)

* Require all channel handlers to be removable
hummingbird-project · Feb 15, 2021 · ed0bb18 · ed0bb18
1 parent 9d04591
commit ed0bb18
Show file tree

Hide file tree

Showing 7 changed files with 28 additions and 13 deletions.
diff --git a/Sources/HummingbirdCore/Server/ChannelInitializer.swift b/Sources/HummingbirdCore/Server/ChannelInitializer.swift
@@ -3,7 +3,12 @@ import NIOHTTP1
 
 /// HTTPServer child channel initializer protocol
 public protocol HBChannelInitializer {
-    func initialize(channel: Channel, childHandlers: [ChannelHandler], configuration: HBHTTPServer.Configuration) -> EventLoopFuture<Void>
+    /// Initialize channel
+    /// - Parameters:
+    ///   - channel: channel
+    ///   - childHandlers: Channel handlers to add
+    ///   - configuration: server configuration
+    func initialize(channel: Channel, childHandlers: [RemovableChannelHandler], configuration: HBHTTPServer.Configuration) -> EventLoopFuture<Void>
 }
 
 /// Setup child channel for HTTP1
@@ -12,10 +17,20 @@ public struct HTTP1ChannelInitializer: HBChannelInitializer {
         self.upgraders = upgraders
     }
 
-    public func initialize(channel: Channel, childHandlers: [ChannelHandler], configuration: HBHTTPServer.Configuration) -> EventLoopFuture<Void> {
+    /// Initialize HTTP1 channel
+    /// - Parameters:
+    ///   - channel: channel
+    ///   - childHandlers: Channel handlers to add
+    ///   - configuration: server configuration
+    public func initialize(channel: Channel, childHandlers: [RemovableChannelHandler], configuration: HBHTTPServer.Configuration) -> EventLoopFuture<Void> {
         var serverUpgrade: NIOHTTPServerUpgradeConfiguration?
         if self.upgraders.count > 0 {
-            serverUpgrade = (self.upgraders, { _ in })
+            serverUpgrade = (self.upgraders, { channel in
+                // remove HTTP handlers after upgrade
+                childHandlers.forEach {
+                    _ = channel.pipeline.removeHandler($0)
+                }
+            })
         }
         return channel.pipeline.configureHTTPServerPipeline(
             withPipeliningAssistance: configuration.withPipeliningAssistance,

diff --git a/Sources/HummingbirdCore/Server/HTTPDecodeHandler.swift b/Sources/HummingbirdCore/Server/HTTPDecodeHandler.swift
@@ -2,7 +2,7 @@ import NIO
 import NIOHTTP1
 
 /// Channel handler for decoding HTTP parts into a HTTP request
-final class HBHTTPDecodeHandler: ChannelDuplexHandler {
+final class HBHTTPDecodeHandler: ChannelDuplexHandler, RemovableChannelHandler {
     public typealias InboundIn = HTTPServerRequestPart
     public typealias InboundOut = HBHTTPRequest
     public typealias OutboundIn = Never

diff --git a/Sources/HummingbirdCore/Server/HTTPEncodeHandler.swift b/Sources/HummingbirdCore/Server/HTTPEncodeHandler.swift
@@ -2,7 +2,7 @@ import NIO
 import NIOHTTP1
 
 /// Channel handler for encoding Response into HTTP parts
-final class HBHTTPEncodeHandler: ChannelOutboundHandler {
+final class HBHTTPEncodeHandler: ChannelOutboundHandler, RemovableChannelHandler {
     typealias OutboundIn = HBHTTPResponse
     typealias OutboundOut = HTTPServerResponsePart
 

diff --git a/Sources/HummingbirdCore/Server/HTTPServer.swift b/Sources/HummingbirdCore/Server/HTTPServer.swift
@@ -47,7 +47,7 @@ public class HBHTTPServer {
     /// - Parameters:
     ///   - handler: autoclosure generating handler
     ///   - position: position to place channel handler
-    @discardableResult public func addChannelHandler(_ handler: @autoclosure @escaping () -> ChannelHandler, position: ChannelPosition = .afterHTTP) -> Self {
+    @discardableResult public func addChannelHandler(_ handler: @autoclosure @escaping () -> RemovableChannelHandler, position: ChannelPosition = .afterHTTP) -> Self {
         self._additionalChildHandlers.append((handler: handler, position: position))
         return self
     }
@@ -129,11 +129,11 @@ public class HBHTTPServer {
         try channel.closeFuture.wait()
     }
 
-    func additionalChannelHandlers(at position: ChannelPosition) -> [ChannelHandler] {
+    func additionalChannelHandlers(at position: ChannelPosition) -> [RemovableChannelHandler] {
         return self._additionalChildHandlers.compactMap { if $0.position == position { return $0.handler() }; return nil }
     }
 
-    private var _additionalChildHandlers: [(handler: () -> ChannelHandler, position: ChannelPosition)]
+    private var _additionalChildHandlers: [(handler: () -> RemovableChannelHandler, position: ChannelPosition)]
 }
 
 extension HBHTTPServer {

diff --git a/Sources/HummingbirdCore/Server/HTTPServerHandler.swift b/Sources/HummingbirdCore/Server/HTTPServerHandler.swift
@@ -3,7 +3,7 @@ import NIO
 import NIOHTTP1
 
 /// Channel handler for responding to a request and returning a response
-final class HBHTTPServerHandler: ChannelInboundHandler {
+final class HBHTTPServerHandler: ChannelInboundHandler, RemovableChannelHandler {
     typealias InboundIn = HBHTTPRequest
     typealias OutboundOut = HBHTTPResponse
 

diff --git a/Sources/HummingbirdHTTP2/ChannelInitializer.swift b/Sources/HummingbirdHTTP2/ChannelInitializer.swift
@@ -5,7 +5,7 @@ import NIOSSL
 
 /// HTTP2 channel initializer
 struct HTTP2ChannelInitializer: HBChannelInitializer {
-    func initialize(channel: Channel, childHandlers: [ChannelHandler], configuration: HBHTTPServer.Configuration) -> EventLoopFuture<Void> {
+    func initialize(channel: Channel, childHandlers: [RemovableChannelHandler], configuration: HBHTTPServer.Configuration) -> EventLoopFuture<Void> {
         return channel.configureHTTP2Pipeline(mode: .server) { streamChannel -> EventLoopFuture<Void> in
             return streamChannel.pipeline.addHandler(HTTP2FramePayloadToHTTP1ServerCodec()).flatMap { () -> EventLoopFuture<Void> in
                 channel.pipeline.addHandlers(childHandlers)
@@ -21,7 +21,7 @@ struct HTTP2UpgradeChannelInitializer: HBChannelInitializer {
     let http1 = HTTP1ChannelInitializer()
     let http2 = HTTP2ChannelInitializer()
 
-    func initialize(channel: Channel, childHandlers: [ChannelHandler], configuration: HBHTTPServer.Configuration) -> EventLoopFuture<Void> {
+    func initialize(channel: Channel, childHandlers: [RemovableChannelHandler], configuration: HBHTTPServer.Configuration) -> EventLoopFuture<Void> {
         channel.configureHTTP2SecureUpgrade(
             h2ChannelConfigurator: { channel in
                 channel.pipeline.addHandlers(childHandlers)

diff --git a/Tests/HummingbirdCoreTests/CoreTests.swift b/Tests/HummingbirdCoreTests/CoreTests.swift
@@ -190,7 +190,7 @@ class HummingBirdCoreTests: XCTestCase {
 
     func testStreamBodySlowStream() {
         /// channel handler that delays the sending of data
-        class SlowInputChannelHandler: ChannelOutboundHandler {
+        class SlowInputChannelHandler: ChannelOutboundHandler, RemovableChannelHandler {
             public typealias OutboundIn = Never
             public typealias OutboundOut = HTTPServerResponsePart
 
@@ -232,7 +232,7 @@ class HummingBirdCoreTests: XCTestCase {
     }
 
     func testChannelHandlerErrorPropagation() {
-        class CreateErrorHandler: ChannelInboundHandler {
+        class CreateErrorHandler: ChannelInboundHandler, RemovableChannelHandler {
             typealias InboundIn = HTTPServerRequestPart
 
             var seen: Bool = false