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

Refactor to add authentication complete flag

Browse Source
This commit is contained in:
David Evans 2021-06-16 15:17:03 +01:00
parent 0252c5d62b
commit 72aa38f3bb
4 changed files with 14 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -76,10 +76,8 @@ public final class SOCKSServerHandshakeHandler: ChannelDuplexHandler, RemovableC
try self.handleWriteSelectedAuthenticationMethod(method, context: context, promise: promise)
case .response(let response):
try self.handleWriteResponse(response, context: context, promise: promise)
case .authenticationData(let data):
try self.handleWriteData(data, context: context, promise: promise)
case .authenticationComplete(let data):
try self.handleAuthenticationComplete(data: data, context: context, promise: promise)
case .authenticationData(let data, let complete):
try self.handleWriteAuthenticationData(data, complete: complete, context: context, promise: promise)
}
} catch {
context.fireErrorCaught(error)
@ -103,18 +101,16 @@ public final class SOCKSServerHandshakeHandler: ChannelDuplexHandler, RemovableC
context.write(self.wrapOutboundOut(buffer), promise: promise)
}
private func handleWriteData(_ data: ByteBuffer, context: ChannelHandlerContext, promise: EventLoopPromise<Void>?) throws {
private func handleWriteAuthenticationData(_ data: ByteBuffer, complete: Bool, context: ChannelHandlerContext, promise: EventLoopPromise<Void>?) throws {
do {
try self.stateMachine.sendData()
if complete {
try self.stateMachine.authenticationComplete()
}
context.write(self.wrapOutboundOut(data), promise: promise)
} catch {
promise?.fail(error)
}
}
private func handleAuthenticationComplete(data: ByteBuffer, context: ChannelHandlerContext, promise: EventLoopPromise<Void>?) throws {
try stateMachine.authenticationComplete()
context.write(self.wrapOutboundOut(data), promise: promise)
}
}

10
Sources/NIOSOCKS/Messages/Messages.swift
View File

@ -38,11 +38,7 @@ public enum ServerMessage: Hashable {
case response(SOCKSResponse)
/// Used when authenticating to send server challenges to the client.
case authenticationData(ByteBuffer)
/// Informs the client that they have been successfully authenticated and
/// can now send the request.
case authenticationComplete(ByteBuffer)
case authenticationData(ByteBuffer, complete: Bool)
}
extension ByteBuffer {
@ -53,10 +49,8 @@ extension ByteBuffer {
return self.writeMethodSelection(method)
case .response(let response):
return self.writeServerResponse(response)
case .authenticationData(var buffer):
case .authenticationData(var buffer, _):
return self.writeBuffer(&buffer)
case .authenticationComplete:
return 0
}
}

2
Tests/NIOSOCKSTests/SOCKSServerHandshakeHandler+Tests+XCTest.swift
View File

@ -30,6 +30,8 @@ extension SOCKSServerHandlerTests {
("testTypicalWorkflowDripfeed", testTypicalWorkflowDripfeed),
("testInboundErrorsAreHandled", testInboundErrorsAreHandled),
("testOutboundErrorsAreHandled", testOutboundErrorsAreHandled),
("testFlushOnHandlerRemoved", testFlushOnHandlerRemoved),
("testForceHandlerRemovalAfterAuth", testForceHandlerRemovalAfterAuth),
]
}
}

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

@ -139,7 +139,7 @@ class SOCKSServerHandlerTests: XCTestCase {
// finish authentication - nothing should be written
// as this is informing the state machine only
self.writeOutbound(.authenticationComplete(ByteBuffer(bytes: [0xFF, 0xFF])))
self.writeOutbound(.authenticationData(ByteBuffer(bytes: [0xFF, 0xFF]), complete: true))
self.assertOutputBuffer([0xFF, 0xFF])
// write the request
@ -178,7 +178,7 @@ class SOCKSServerHandlerTests: XCTestCase {
// finish authentication - nothing should be written
// as this is informing the state machine only
XCTAssertNoThrow(try self.channel.writeOutbound(ServerMessage.authenticationComplete(ByteBuffer(bytes: [0xFF, 0xFF]))))
XCTAssertNoThrow(try self.channel.writeOutbound(ServerMessage.authenticationData(ByteBuffer(bytes: [0xFF, 0xFF]), complete: true)))
self.assertOutputBuffer([0xFF, 0xFF])
// write the request
@ -202,7 +202,7 @@ class SOCKSServerHandlerTests: XCTestCase {
// write something that will be be invalid for the state machine's
// current state, causing an error to be thrown
func testOutboundErrorsAreHandled() {
XCTAssertThrowsError(try self.channel.writeAndFlush(ServerMessage.authenticationComplete(ByteBuffer(bytes: [0xFF, 0xFF]))).wait()) { e in
XCTAssertThrowsError(try self.channel.writeAndFlush(ServerMessage.authenticationData(ByteBuffer(bytes: [0xFF, 0xFF]), complete: true)).wait()) { e in
XCTAssertTrue(e is SOCKSError.InvalidServerState)
}
}
@ -227,6 +227,6 @@ class SOCKSServerHandlerTests: XCTestCase {
// auth complete, try to write data without
// removing the handler, it should fail
XCTAssertThrowsError(try self.channel.writeOutbound(ServerMessage.authenticationData(ByteBuffer(string: "hello, world!"))))
XCTAssertThrowsError(try self.channel.writeOutbound(ServerMessage.authenticationData(ByteBuffer(string: "hello, world!"), complete: false)))
}
}