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

Add data to authentication complete

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

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

4
Sources/NIOSOCKS/Messages/Errors.swift
View File

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

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

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

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

@ -139,8 +139,8 @@ 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) self.writeOutbound(.authenticationComplete(ByteBuffer(bytes: [0xFF, 0xFF])))
self.assertOutputBuffer([]) self.assertOutputBuffer([0xFF, 0xFF])
// write the request // write the request
XCTAssertFalse(testHandler.hadRequest) XCTAssertFalse(testHandler.hadRequest)
@ -178,8 +178,8 @@ 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)) XCTAssertNoThrow(try self.channel.writeOutbound(ServerMessage.authenticationComplete(ByteBuffer(bytes: [0xFF, 0xFF]))))
self.assertOutputBuffer([]) self.assertOutputBuffer([0xFF, 0xFF])
// write the request // write the request
XCTAssertFalse(testHandler.hadRequest) XCTAssertFalse(testHandler.hadRequest)
@ -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).wait()) { e in XCTAssertThrowsError(try self.channel.writeAndFlush(ServerMessage.authenticationComplete(ByteBuffer(bytes: [0xFF, 0xFF]))).wait()) { e in
XCTAssertTrue(e is SOCKSError.InvalidServerState) XCTAssertTrue(e is SOCKSError.InvalidServerState)
} }
} }