apple/swift-nio-extras
1
0
Fork 0
mirror of https://github.com/apple/swift-nio-extras.git synced 2025-05-14 17:02:43 +08:00
Code Issues Packages Projects Releases Wiki Activity

Make handler removable

Browse Source
This commit is contained in:
David Evans 2021-06-16 13:23:55 +01:00
parent 13fda5c503
commit 0a780c8c72
2 changed files with 20 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
Sources/NIOSOCKS/Channel Handlers/SOCKSServerHandshakeHandler.swift
View File

@ -19,7 +19,7 @@ import NIO
/// and parser to enforce SOCKSv5 protocol correctness. Inbound bytes will by parsed into /// and parser to enforce SOCKSv5 protocol correctness. Inbound bytes will by parsed into
/// `ClientMessage` for downstream consumption. Send `ServerMessage` to this /// `ClientMessage` for downstream consumption. Send `ServerMessage` to this
/// handler. /// handler.
public final class SOCKSServerHandshakeHandler: ChannelDuplexHandler { public final class SOCKSServerHandshakeHandler: ChannelDuplexHandler, RemovableChannelHandler {
public typealias InboundIn = ByteBuffer public typealias InboundIn = ByteBuffer
public typealias InboundOut = ClientMessage public typealias InboundOut = ClientMessage
@ -61,6 +61,13 @@ public final class SOCKSServerHandshakeHandler: ChannelDuplexHandler {
} }
} }
public func handlerRemoved(context: ChannelHandlerContext) {
guard let buffer = self.inboundBuffer else {
return
}
context.fireChannelRead(.init(buffer))
}
public func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) { public func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
do { do {
let message = self.unwrapOutboundIn(data) let message = self.unwrapOutboundIn(data)

14
Tests/NIOSOCKSTests/SOCKSServerHandshakeHandler+Tests.swift
View File

@ -91,8 +91,11 @@ class SOCKSServerHandlerTests: XCTestCase {
} }
func assertInbound(_ bytes: [UInt8], line: UInt = #line) { func assertInbound(_ bytes: [UInt8], line: UInt = #line) {
var buffer = try! self.channel.readInbound(as: ByteBuffer.self) if var buffer = try! self.channel.readInbound(as: ByteBuffer.self) {
XCTAssertEqual(buffer!.readBytes(length: buffer!.readableBytes), bytes, line: line) XCTAssertEqual(buffer.readBytes(length: buffer.readableBytes), bytes, line: line)
} else {
XCTAssertTrue(bytes.count == 0)
}
} }
func testTypicalWorkflow() { func testTypicalWorkflow() {
@ -198,4 +201,11 @@ class SOCKSServerHandlerTests: XCTestCase {
XCTAssertTrue(e is SOCKSError.InvalidServerState) XCTAssertTrue(e is SOCKSError.InvalidServerState)
} }
} }
func testFlushOnHandlerRemoved() {
self.writeInbound([0x05, 0x01])
self.assertInbound([])
XCTAssertNoThrow(try self.channel.pipeline.removeHandler(self.handler).wait())
self.assertInbound([0x05, 0x01])
}
} }