Add data to authentication complete

This commit is contained in:
David Evans 2021-06-16 15:11:53 +01:00
parent 2c8c151ae8
commit 0252c5d62b
4 changed files with 22 additions and 18 deletions

View File

@ -35,13 +35,13 @@ public final class SOCKSServerHandshakeHandler: ChannelDuplexHandler, RemovableC
public func channelRead(context: ChannelHandlerContext, data: NIOAny) {
var message = self.unwrapInboundIn(data)
self.inboundBuffer.setOrWriteBuffer(&message)
if self.stateMachine.proxyEstablished {
context.fireChannelRead(data)
return
}
var message = self.unwrapInboundIn(data)
self.inboundBuffer.setOrWriteBuffer(&message)
do {
// safe to bang inbound buffer, it's always written above
guard let message = try self.stateMachine.receiveBuffer(&self.inboundBuffer!) else {
@ -78,8 +78,8 @@ public final class SOCKSServerHandshakeHandler: ChannelDuplexHandler, RemovableC
try self.handleWriteResponse(response, context: context, promise: promise)
case .authenticationData(let data):
try self.handleWriteData(data, context: context, promise: promise)
case .authenticationComplete:
try self.handleAuthenticationComplete(context: context, promise: promise)
case .authenticationComplete(let data):
try self.handleAuthenticationComplete(data: data, context: context, promise: promise)
}
} catch {
context.fireErrorCaught(error)
@ -103,7 +103,7 @@ 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 handleWriteData(_ data: ByteBuffer, context: ChannelHandlerContext, promise: EventLoopPromise<Void>?) throws {
do {
try self.stateMachine.sendData()
context.write(self.wrapOutboundOut(data), promise: promise)
@ -112,9 +112,9 @@ public final class SOCKSServerHandshakeHandler: ChannelDuplexHandler, RemovableC
}
}
private func handleAuthenticationComplete(context: ChannelHandlerContext, promise: EventLoopPromise<Void>?) throws {
private func handleAuthenticationComplete(data: ByteBuffer, context: ChannelHandlerContext, promise: EventLoopPromise<Void>?) throws {
try stateMachine.authenticationComplete()
promise?.succeed(())
context.write(self.wrapOutboundOut(data), promise: promise)
}
}

View File

@ -19,12 +19,16 @@ public enum SOCKSError {
/// The SOCKS client was in a different state to that required.
public struct InvalidClientState: Error, Hashable {
public init() {
}
}
/// The SOCKS server was in a different state to that required.
public struct InvalidServerState: Error, Hashable {
public init() {
}
}
/// The protocol version was something other than *5*. Note that

View File

@ -40,9 +40,9 @@ public enum ServerMessage: Hashable {
/// Used when authenticating to send server challenges to the client.
case authenticationData(ByteBuffer)
/// This is a faux message to update the server's state machine. It should be sent
/// once the server is satisified that the client is fully-authenticated.
case authenticationComplete
/// Informs the client that they have been successfully authenticated and
/// can now send the request.
case authenticationComplete(ByteBuffer)
}
extension ByteBuffer {

View File

@ -139,8 +139,8 @@ class SOCKSServerHandlerTests: XCTestCase {
// finish authentication - nothing should be written
// as this is informing the state machine only
self.writeOutbound(.authenticationComplete)
self.assertOutputBuffer([])
self.writeOutbound(.authenticationComplete(ByteBuffer(bytes: [0xFF, 0xFF])))
self.assertOutputBuffer([0xFF, 0xFF])
// write the request
XCTAssertFalse(testHandler.hadRequest)
@ -178,8 +178,8 @@ class SOCKSServerHandlerTests: XCTestCase {
// finish authentication - nothing should be written
// as this is informing the state machine only
XCTAssertNoThrow(try self.channel.writeOutbound(ServerMessage.authenticationComplete))
self.assertOutputBuffer([])
XCTAssertNoThrow(try self.channel.writeOutbound(ServerMessage.authenticationComplete(ByteBuffer(bytes: [0xFF, 0xFF]))))
self.assertOutputBuffer([0xFF, 0xFF])
// write the request
XCTAssertFalse(testHandler.hadRequest)
@ -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).wait()) { e in
XCTAssertThrowsError(try self.channel.writeAndFlush(ServerMessage.authenticationComplete(ByteBuffer(bytes: [0xFF, 0xFF]))).wait()) { e in
XCTAssertTrue(e is SOCKSError.InvalidServerState)
}
}