apple/swift-nio-extras
1
0
Fork 0
mirror of https://github.com/apple/swift-nio-extras.git synced 2025-05-15 09:22:38 +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) try self.handleWriteSelectedAuthenticationMethod(method, context: context, promise: promise)
case .response(let response): case .response(let response):
try self.handleWriteResponse(response, context: context, promise: promise) try self.handleWriteResponse(response, context: context, promise: promise)
case .authenticationData(let data): case .authenticationData(let data, let complete):
try self.handleWriteData(data, context: context, promise: promise) try self.handleWriteAuthenticationData(data, complete: complete, context: context, promise: promise)
case .authenticationComplete(let data):
try self.handleAuthenticationComplete(data: data, context: context, promise: promise)
} }
} catch { } catch {
context.fireErrorCaught(error) context.fireErrorCaught(error)
@ -103,18 +101,16 @@ public final class SOCKSServerHandshakeHandler: ChannelDuplexHandler, RemovableC
context.write(self.wrapOutboundOut(buffer), promise: promise) 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 { do {
try self.stateMachine.sendData() try self.stateMachine.sendData()
if complete {
try self.stateMachine.authenticationComplete()
}
context.write(self.wrapOutboundOut(data), promise: promise) context.write(self.wrapOutboundOut(data), promise: promise)
} catch { } catch {
promise?.fail(error) 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) case response(SOCKSResponse)
/// Used when authenticating to send server challenges to the client. /// Used when authenticating to send server challenges to the client.
case authenticationData(ByteBuffer) case authenticationData(ByteBuffer, complete: Bool)
/// Informs the client that they have been successfully authenticated and
/// can now send the request.
case authenticationComplete(ByteBuffer)
} }
extension ByteBuffer { extension ByteBuffer {
@ -53,10 +49,8 @@ extension ByteBuffer {
return self.writeMethodSelection(method) return self.writeMethodSelection(method)
case .response(let response): case .response(let response):
return self.writeServerResponse(response) return self.writeServerResponse(response)
case .authenticationData(var buffer): case .authenticationData(var buffer, _):
return self.writeBuffer(&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), ("testTypicalWorkflowDripfeed", testTypicalWorkflowDripfeed),
("testInboundErrorsAreHandled", testInboundErrorsAreHandled), ("testInboundErrorsAreHandled", testInboundErrorsAreHandled),
("testOutboundErrorsAreHandled", testOutboundErrorsAreHandled), ("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 // finish authentication - nothing should be written
// as this is informing the state machine only // 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]) self.assertOutputBuffer([0xFF, 0xFF])
// write the request // write the request
@ -178,7 +178,7 @@ class SOCKSServerHandlerTests: XCTestCase {
// finish authentication - nothing should be written // finish authentication - nothing should be written
// as this is informing the state machine only // 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]) self.assertOutputBuffer([0xFF, 0xFF])
// write the request // write the request
@ -202,7 +202,7 @@ class SOCKSServerHandlerTests: XCTestCase {
// write something that will be be invalid for the state machine's // write something that will be be invalid for the state machine's
// current state, causing an error to be thrown // current state, causing an error to be thrown
func testOutboundErrorsAreHandled() { 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) XCTAssertTrue(e is SOCKSError.InvalidServerState)
} }
} }
@ -227,6 +227,6 @@ class SOCKSServerHandlerTests: XCTestCase {
// auth complete, try to write data without // auth complete, try to write data without
// removing the handler, it should fail // 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)))
} }
} }