punchnet-macos/Tun/PacketTunnelProvider.swift

//
//  PacketTunnelProvider.swift
//  punchnet
//
//  Created by 安礼成 on 2025/8/3.
//

import Foundation
import NetworkExtension

enum TunnelError: Error {
    case invalidConfiguration
    case invalidContext
}

class PacketTunnelProvider: NEPacketTunnelProvider {
    private enum RuntimeState {
        case idle
        case starting
        case running
        case stopping
    }

    private let runtimeLock = NSLock()
    private var runtimeState: RuntimeState = .idle
    private var config: SDLConfiguration?
    private var rsaCipher: CCRSACipher?
    private var contextActor: SDLContextActor?
    private var startCompletionHandler: ((Error?) -> Void)?
    
    override func startTunnel(options: [String: NSObject]?, completionHandler: @escaping (Error?) -> Void) {
        guard let options, let config = SDLConfiguration.parse(options: options) else {
            completionHandler(TunnelError.invalidConfiguration)
            return
        }

        let rsaCipher = try! CCRSACipher(keySize: 1024)
        self.startContext(config: config, rsaCipher: rsaCipher, completionHandler: completionHandler)
    }

    override func stopTunnel(with reason: NEProviderStopReason, completionHandler: @escaping () -> Void) {
        self.stopContext(clearRuntimeConfiguration: true, completionHandler: completionHandler)
    }

    override func handleAppMessage(_ messageData: Data, completionHandler: ((Data?) -> Void)?) {
        // Add code here to handle the message.
        Task {
            do {
                let message = try AppRequest(serializedBytes: messageData)
                let replyData = try await self.handleAppRequest(message: message)
                completionHandler?(replyData)
            } catch let err {
                var reply = TunnelResponse()
                reply.code = 1
                reply.message = err.localizedDescription

                let errorReplyData = try? reply.serializedData()
                completionHandler?(errorReplyData)
            }
        }
    }

    override func sleep(completionHandler: @escaping () -> Void) {
        self.stopContext(clearRuntimeConfiguration: false) {
            SDLLogger.log("[PacketTunnelProvider] sleep")
            completionHandler()
        }
    }

    override func wake() {
        SDLLogger.log("[PacketTunnelProvider] wake up!!!!!!!")
        self.startCachedContext { err in
            if let err {
                SDLLogger.log("[PacketTunnelProvider] wakeup start failed: \(err.localizedDescription)")
            } else {
                SDLLogger.log("[PacketTunnelProvider] wakeup and try start")
            }
        }
    }

    private func handleAppRequest(message: AppRequest) async throws -> Data? {
        guard let contextActor = self.currentContextActor() else {
            throw TunnelError.invalidContext
        }
        
        switch message.command {
        case .changeExitNode(let changeExitNode):
            let exitNodeIp = changeExitNode.ip
            do {
                try await contextActor.updateExitNode(exitNodeIp: exitNodeIp)
                var reply = TunnelResponse()
                reply.code = 0
                reply.message = "操作成功"
                return try reply.serializedData()
                
            } catch let err {
                var reply = TunnelResponse()
                reply.code = 1
                reply.message = err.localizedDescription
                
                return try reply.serializedData()
            }
        case .none:
            var reply = TunnelResponse()
            reply.code = 1
            reply.message = "无效请求"
            return try reply.serializedData()
        }
    }

}

private extension PacketTunnelProvider {

    func startCachedContext(completionHandler: @escaping (Error?) -> Void) {
        self.runtimeLock.lock()
        let config = self.config
        let rsaCipher = self.rsaCipher
        self.runtimeLock.unlock()

        guard let config, let rsaCipher else {
            completionHandler(TunnelError.invalidConfiguration)
            return
        }

        self.startContext(config: config, rsaCipher: rsaCipher, completionHandler: completionHandler)
    }

    func startContext(config: SDLConfiguration, rsaCipher: CCRSACipher, completionHandler: @escaping (Error?) -> Void) {
        self.runtimeLock.lock()
        switch self.runtimeState {
        case .idle:
            SDLTunnelAppNotifier.shared.clear()

            let contextActor = SDLContextActor(
                provider: self,
                config: config,
                rsaCipher: rsaCipher
            )

            self.config = config
            self.rsaCipher = rsaCipher
            self.contextActor = contextActor
            self.startCompletionHandler = completionHandler
            self.runtimeState = .starting
            self.runtimeLock.unlock()

            Task {
                do {
                    try await contextActor.start()
                    self.finishContextStart(contextActor, error: nil)
                } catch {
                    self.finishContextStart(contextActor, error: error)
                }
            }

        case .starting, .running, .stopping:
            self.runtimeLock.unlock()
            completionHandler(TunnelError.invalidContext)
        }
    }

    func stopContext(clearRuntimeConfiguration: Bool, completionHandler: @escaping () -> Void) {
        self.runtimeLock.lock()
        let contextActor = self.contextActor
        let startCompletionHandler = self.startCompletionHandler

        guard contextActor != nil else {
            self.runtimeState = .idle
            self.startCompletionHandler = nil
            if clearRuntimeConfiguration {
                self.config = nil
                self.rsaCipher = nil
            }
            self.runtimeLock.unlock()
            startCompletionHandler?(TunnelError.invalidContext)
            completionHandler()
            return
        }

        self.contextActor = nil
        self.startCompletionHandler = nil
        self.runtimeState = .stopping
        if clearRuntimeConfiguration {
            self.config = nil
            self.rsaCipher = nil
        }
        self.runtimeLock.unlock()

        startCompletionHandler?(TunnelError.invalidContext)

        Task {
            await contextActor?.stop()
            self.markContextStopped()
            completionHandler()
        }
    }

    func currentContextActor() -> SDLContextActor? {
        self.runtimeLock.lock()
        let contextActor = self.contextActor
        self.runtimeLock.unlock()
        return contextActor
    }

    func finishContextStart(_ contextActor: SDLContextActor, error: Error?) {
        self.runtimeLock.lock()
        guard self.contextActor === contextActor else {
            self.runtimeLock.unlock()
            return
        }

        let startCompletionHandler = self.startCompletionHandler
        self.startCompletionHandler = nil

        if let error {
            self.contextActor = nil
            self.runtimeState = .idle
        } else {
            self.runtimeState = .running
        }
        self.runtimeLock.unlock()

        startCompletionHandler?(error)
    }

    func markContextStopped() {
        self.runtimeLock.lock()
        if self.contextActor == nil {
            self.runtimeState = .idle
        }
        self.runtimeLock.unlock()
    }
}