ChannelPipeline: fail double removal of handler (#1133)

Sources/NIO/Channel.swift

@@ -356,6 +356,9 @@ extension ChannelError: Equatable { }
 /// we have seen this happen on Darwin.
 public struct NIOFailedToSetSocketNonBlockingError: Error {}
 
+/// The removal of a `ChannelHandler` using `ChannelPipeline.removeHandler` has been attempted more than once.
+public struct NIOAttemptedToRemoveHandlerMultipleTimesError: Error {}
+
 /// An `Channel` related event that is passed through the `ChannelPipeline` to notify the user.
 public enum ChannelEvent: Equatable {
     /// `ChannelOptions.allowRemoteHalfClosure` is `true` and input portion of the `Channel` was closed.

Sources/NIO/ChannelPipeline.swift

@@ -400,11 +400,15 @@ public final class ChannelPipeline: ChannelInvoker {
             promise?.fail(ChannelError.unremovableHandler)
             return
         }
+        func removeHandler0() {
+            context.startUserTriggeredRemoval(promise: promise)
+        }
+
         if self.eventLoop.inEventLoop {
-            handler.removeHandler(context: context, removalToken: .init(promise: promise))
+            removeHandler0()
         } else {
             self.eventLoop.execute {
-                handler.removeHandler(context: context, removalToken: .init(promise: promise))
+                removeHandler0()
             }
         }
     }
@@ -1100,6 +1104,7 @@ public final class ChannelHandlerContext: ChannelInvoker {
     private let inboundHandler: _ChannelInboundHandler?
     private let outboundHandler: _ChannelOutboundHandler?
     private var removeHandlerInvoked = false
+    private var userTriggeredRemovalStarted = false
 
     // Only created from within ChannelPipeline
     fileprivate init(name: String, handler: ChannelHandler, pipeline: ChannelPipeline) {
@@ -1507,6 +1512,17 @@ extension ChannelHandlerContext {
         self.eventLoop.preconditionInEventLoop()
         self.pipeline.removeHandlerFromPipeline(context: self, promise: removalToken.promise)
     }
+
+    internal func startUserTriggeredRemoval(promise: EventLoopPromise<Void>?) {
+        self.eventLoop.assertInEventLoop()
+        guard !self.userTriggeredRemovalStarted else {
+            promise?.fail(NIOAttemptedToRemoveHandlerMultipleTimesError())
+            return
+        }
+        self.userTriggeredRemovalStarted = true
+        (self.handler as! RemovableChannelHandler).removeHandler(context: self,
+                                                                 removalToken: .init(promise: promise))
+    }
 }
 
 extension ChannelPipeline: CustomDebugStringConvertible {

Tests/NIOTests/ChannelPipelineTest+XCTest.swift

@@ -61,6 +61,7 @@ extension ChannelPipelineTest {
                 ("testAddMultipleHandlers", testAddMultipleHandlers),
                 ("testPipelineDebugDescription", testPipelineDebugDescription),
                 ("testWeDontCallHandlerRemovedTwiceIfAHandlerCompletesRemovalOnlyAfterChannelTeardown", testWeDontCallHandlerRemovedTwiceIfAHandlerCompletesRemovalOnlyAfterChannelTeardown),
+                ("testWeFailTheSecondRemoval", testWeFailTheSecondRemoval),
            ]
    }
 }

Tests/NIOTests/ChannelPipelineTest.swift

@@ -1288,6 +1288,50 @@ class ChannelPipelineTest: XCTestCase {
 
         XCTAssertNoThrow(try allDonePromise.futureResult.wait())
     }
+
+    func testWeFailTheSecondRemoval() {
+        final class Handler: ChannelInboundHandler, RemovableChannelHandler {
+            typealias InboundIn = Never
+
+            private let removalTriggeredPromise: EventLoopPromise<Void>
+            private let continueRemovalFuture: EventLoopFuture<Void>
+            private var removeHandlerCalls = 0
+
+            init(removalTriggeredPromise: EventLoopPromise<Void>,
+                 continueRemovalFuture: EventLoopFuture<Void>) {
+                self.removalTriggeredPromise = removalTriggeredPromise
+                self.continueRemovalFuture = continueRemovalFuture
+            }
+
+            func removeHandler(context: ChannelHandlerContext, removalToken: ChannelHandlerContext.RemovalToken) {
+                self.removeHandlerCalls += 1
+                XCTAssertEqual(1, self.removeHandlerCalls)
+                self.removalTriggeredPromise.succeed(())
+                self.continueRemovalFuture.whenSuccess {
+                    context.leavePipeline(removalToken: removalToken)
+                }
+            }
+        }
+
+        let channel = EmbeddedChannel()
+        let removalTriggeredPromise: EventLoopPromise<Void> = channel.eventLoop.makePromise()
+        let continueRemovalPromise: EventLoopPromise<Void> = channel.eventLoop.makePromise()
+
+        let handler = Handler(removalTriggeredPromise: removalTriggeredPromise,
+                              continueRemovalFuture: continueRemovalPromise.futureResult)
+        XCTAssertNoThrow(try channel.pipeline.addHandler(handler).wait())
+        let removal1Future = channel.pipeline.removeHandler(handler)
+        XCTAssertThrowsError(try channel.pipeline.removeHandler(handler).wait()) { error in
+            XCTAssert(error is NIOAttemptedToRemoveHandlerMultipleTimesError,
+                      "unexpected error: \(error)")
+        }
+        continueRemovalPromise.succeed(())
+        XCTAssertThrowsError(try channel.pipeline.removeHandler(handler).wait()) { error in
+            XCTAssertEqual(.notFound, error as? ChannelPipelineError, "unexpected error: \(error)")
+        }
+        XCTAssertNoThrow(try removal1Future.wait())
+        XCTAssertNoThrow(XCTAssertTrue(try channel.finish().isClean))
+    }
 }
 
 // this should be within `testAddMultipleHandlers` but https://bugs.swift.org/browse/SR-9956