Skip to content

gist: fix socks5 udp #90

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
DmitriyMV wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

gist: fix socks5 udp #90

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 53 additions & 13 deletions MacOS/ProxyBridge/extension/AppProxyProvider.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -339,11 +339,37 @@ class AppProxyProvider: NETransparentProxyProvider {
}

override func stopProxy(with reason: NEProviderStopReason, completionHandler: @escaping () -> Void) {
tcpConnectionsLock.lock()
let drainedUDPResources = udpResources
udpResources.removeAll()
tcpConnectionsLock.unlock()

for (flow, resources) in drainedUDPResources {
resources.udpSession?.cancel()
resources.tcpConnection.cancel()
flow.closeReadWithError(nil)
flow.closeWriteWithError(nil)
}

completionHandler()
}

private var udpTCPConnections: [NEAppProxyUDPFlow: NWTCPConnection] = [:]
private var udpResources: [NEAppProxyUDPFlow: (tcpConnection: NWTCPConnection, udpSession: NWUDPSession)] = [:]
private let tcpConnectionsLock = NSLock()
private let udpRetryQueue = DispatchQueue(label: "com.interceptsuite.ProxyBridge.udp-retry", qos: .utility)

private func cleanupUDPResources(for clientFlow: NEAppProxyUDPFlow, error: Error? = nil) {
tcpConnectionsLock.lock()
let removedResources = udpResources.removeValue(forKey: clientFlow)
tcpConnectionsLock.unlock()

if let (tcpConnection, udpSession) = removedResources {
udpSession.cancel()
tcpConnection.cancel()
}
clientFlow.closeReadWithError(error)
clientFlow.closeWriteWithError(error)
}

override func handleAppMessage(_ messageData: Data, completionHandler: ((Data?) -> Void)?) {
guard let message = try? JSONSerialization.jsonObject(with: messageData) as? [String: Any],
Expand Down Expand Up @@ -638,6 +664,7 @@ class AppProxyProvider: NETransparentProxyProvider {

if let error = error {
self.log("Failed to open UDP flow: \(error.localizedDescription)", level: "ERROR")
self.cleanupUDPResources(for: flow, error: error)
return
}

Expand All @@ -664,20 +691,20 @@ class AppProxyProvider: NETransparentProxyProvider {
guard let self = self else { return }

if let error = error {
self.log("SOCKS5 UDP greeting failed: \(error.localizedDescription)", level: "ERROR")
self.failUDPSetup(for: clientFlow, tcpConnection: tcpConnection, message: "SOCKS5 UDP greeting failed: \(error.localizedDescription)", error: error)
return
}

tcpConnection.readMinimumLength(2, maximumLength: 2) { [weak self] data, error in
guard let self = self else { return }

if let error = error {
self.log("SOCKS5 UDP greeting response failed: \(error.localizedDescription)", level: "ERROR")
self.failUDPSetup(for: clientFlow, tcpConnection: tcpConnection, message: "SOCKS5 UDP greeting response failed: \(error.localizedDescription)", error: error)
return
}

guard let data = data, data.count == 2, data[1] == 0x00 else {
self.log("SOCKS5 UDP greeting response failed", level: "ERROR")
self.failUDPSetup(for: clientFlow, tcpConnection: tcpConnection, message: "SOCKS5 UDP greeting response failed")
return
}

Expand All @@ -699,20 +726,20 @@ class AppProxyProvider: NETransparentProxyProvider {
guard let self = self else { return }

if let error = error {
self.log("SOCKS5 UDP ASSOCIATE failed: \(error.localizedDescription)", level: "ERROR")
self.failUDPSetup(for: clientFlow, tcpConnection: tcpConnection, message: "SOCKS5 UDP ASSOCIATE failed: \(error.localizedDescription)", error: error)
return
}

tcpConnection.readMinimumLength(10, maximumLength: 512) { [weak self] data, error in
guard let self = self else { return }

if let error = error {
self.log("SOCKS5 UDP ASSOCIATE response error: \(error.localizedDescription)", level: "ERROR")
self.failUDPSetup(for: clientFlow, tcpConnection: tcpConnection, message: "SOCKS5 UDP ASSOCIATE response error: \(error.localizedDescription)", error: error)
return
}

guard let data = data, data.count >= 10, data[0] == 0x05, data[1] == 0x00 else {
self.log("SOCKS5 UDP ASSOCIATE rejected", level: "ERROR")
self.failUDPSetup(for: clientFlow, tcpConnection: tcpConnection, message: "SOCKS5 UDP ASSOCIATE rejected")
return
}

Expand All @@ -722,6 +749,12 @@ class AppProxyProvider: NETransparentProxyProvider {
}
}

private func failUDPSetup(for clientFlow: NEAppProxyUDPFlow, tcpConnection: NWTCPConnection, message: String, error: Error? = nil) {
log(message, level: "ERROR")
cleanupUDPResources(for: clientFlow, error: error)
tcpConnection.cancel()
}

private func parseSOCKS5Address(from data: Data, offset: Int) -> (String, UInt16) {
let atyp = data[offset]

Expand Down Expand Up @@ -754,26 +787,32 @@ class AppProxyProvider: NETransparentProxyProvider {
let udpSession = self.createUDPSession(to: relayEndpoint, from: nil)

tcpConnectionsLock.lock()
udpTCPConnections[clientFlow] = tcpConnection
udpResources[clientFlow] = (tcpConnection: tcpConnection, udpSession: udpSession)
tcpConnectionsLock.unlock()

readAndForwardClientUDP(clientFlow: clientFlow, udpSession: udpSession, processPath: processPath)
readAndForwardRelayUDP(clientFlow: clientFlow, udpSession: udpSession)
}

private func readAndForwardClientUDP(clientFlow: NEAppProxyUDPFlow, udpSession: NWUDPSession, processPath: String) {
private func readAndForwardClientUDP(clientFlow: NEAppProxyUDPFlow, udpSession: NWUDPSession, processPath: String, emptyReadCount: Int = 0) {
var isFirstPacket = true

clientFlow.readDatagrams { [weak self] datagrams, endpoints, error in
guard let self = self else { return }

if let error = error {
self.log("UDP read error: \(error.localizedDescription)", level: "ERROR")
self.cleanupUDPResources(for: clientFlow, error: error)
return
}

guard let datagrams = datagrams, let endpoints = endpoints, !datagrams.isEmpty else {
self.readAndForwardClientUDP(clientFlow: clientFlow, udpSession: udpSession, processPath: processPath)
let cappedEmptyReadCount = min(emptyReadCount + 1, 4)
let delayMs = 50 * (1 << min(emptyReadCount, 4))

self.udpRetryQueue.asyncAfter(deadline: .now() + .milliseconds(delayMs)) {
self.readAndForwardClientUDP(clientFlow: clientFlow, udpSession: udpSession, processPath: processPath, emptyReadCount: cappedEmptyReadCount)
}
return
}

Expand All @@ -800,12 +839,13 @@ class AppProxyProvider: NETransparentProxyProvider {
udpSession.writeDatagram(encapsulated, completionHandler: { error in
if let error = error {
self.log("UDP write error: \(error)", level: "ERROR")
self.cleanupUDPResources(for: clientFlow, error: error)
}
})
}
}

self.readAndForwardClientUDP(clientFlow: clientFlow, udpSession: udpSession, processPath: processPath)
self.readAndForwardClientUDP(clientFlow: clientFlow, udpSession: udpSession, processPath: processPath, emptyReadCount: 0)
}
}

Expand All @@ -815,6 +855,7 @@ class AppProxyProvider: NETransparentProxyProvider {

if let error = error {
self.log("UDP relay error: \(error.localizedDescription)", level: "ERROR")
self.cleanupUDPResources(for: clientFlow, error: error)
return
}

Expand All @@ -833,6 +874,7 @@ class AppProxyProvider: NETransparentProxyProvider {
clientFlow.writeDatagrams(unwrappedDatagrams, sentBy: unwrappedEndpoints) { error in
if let error = error {
self.log("UDP response write error: \(error.localizedDescription)", level: "ERROR")
self.cleanupUDPResources(for: clientFlow, error: error)
}
}
}
Expand Down Expand Up @@ -1355,5 +1397,3 @@ class AppProxyProvider: NETransparentProxyProvider {
logQueueLock.unlock()
}
}