94 lines
3.8 KiB
Swift
94 lines
3.8 KiB
Swift
import Foundation
|
|
import Network
|
|
import UltraMeshCore
|
|
|
|
final class PeerSession {
|
|
let connection: NWConnection
|
|
let outbound: Bool
|
|
var remote: NodeRecord?
|
|
var onMessage: ((PeerSession, WireMessage) -> Void)?
|
|
var onStop: ((PeerSession) -> Void)?
|
|
private let queue: DispatchQueue
|
|
private var stopped = false
|
|
private struct Pending { let data: Data; let bulk: Bool }
|
|
private var priorityQueue: [Pending] = []
|
|
private var bulkQueue: [Pending] = []
|
|
private var bulkBytes = 0
|
|
private var sending = false
|
|
private static let maximumPendingBulkBytes = 512 * 1024
|
|
|
|
init(connection: NWConnection, outbound: Bool, queue: DispatchQueue) {
|
|
self.connection = connection; self.outbound = outbound; self.queue = queue
|
|
}
|
|
|
|
func start(local: NodeRecord) {
|
|
connection.stateUpdateHandler = { [weak self] state in
|
|
guard let self else { return }
|
|
switch state {
|
|
case .ready:
|
|
self.send(.hello(local)); self.receiveMessage()
|
|
case .failed, .cancelled: self.stop()
|
|
default: break
|
|
}
|
|
}
|
|
connection.start(queue: queue)
|
|
}
|
|
|
|
func send(_ message: WireMessage) {
|
|
guard !stopped, let data = try? FrameCodec.encode(WireEnvelope(message: message)) else { return }
|
|
let bulk: Bool
|
|
if case .packet(let packet) = message { bulk = packet.trafficClass == .nativeIPv6 } else { bulk = false }
|
|
if bulk {
|
|
guard bulkBytes + data.count <= Self.maximumPendingBulkBytes else { return }
|
|
bulkBytes += data.count; bulkQueue.append(.init(data: data, bulk: true))
|
|
} else { priorityQueue.append(.init(data: data, bulk: false)) }
|
|
pumpWrites()
|
|
}
|
|
|
|
func stop() {
|
|
guard !stopped else { return }; stopped = true
|
|
connection.cancel(); onStop?(self)
|
|
}
|
|
|
|
private func receiveMessage() {
|
|
receiveExactly(4, accumulated: Data()) { [weak self] prefix in
|
|
guard let self, let prefix,
|
|
let size = try? FrameCodec.bodyLength(from: prefix) else { self?.stop(); return }
|
|
self.receiveExactly(size, accumulated: Data()) { [weak self] body in
|
|
guard let self, let body,
|
|
let envelope = try? FrameCodec.decode(WireEnvelope.self, body: body),
|
|
envelope.version == WireEnvelope.currentVersion else { self?.stop(); return }
|
|
self.onMessage?(self, envelope.message)
|
|
self.receiveMessage()
|
|
}
|
|
}
|
|
}
|
|
|
|
private func pumpWrites() {
|
|
guard !stopped, !sending else { return }
|
|
let pending: Pending
|
|
if !priorityQueue.isEmpty { pending = priorityQueue.removeFirst() }
|
|
else if !bulkQueue.isEmpty { pending = bulkQueue.removeFirst(); bulkBytes -= pending.data.count }
|
|
else { return }
|
|
sending = true
|
|
connection.send(content: pending.data, completion: .contentProcessed { [weak self] error in
|
|
guard let self else { return }
|
|
self.queue.async {
|
|
self.sending = false
|
|
if error != nil { self.stop() } else { self.pumpWrites() }
|
|
}
|
|
})
|
|
}
|
|
|
|
private func receiveExactly(_ count: Int, accumulated: Data, completion: @escaping (Data?) -> Void) {
|
|
if accumulated.count == count { completion(accumulated); return }
|
|
connection.receive(minimumIncompleteLength: 1, maximumLength: count - accumulated.count) { [weak self] data, _, complete, error in
|
|
guard self != nil, error == nil, let data, !data.isEmpty else { completion(nil); return }
|
|
let joined = accumulated + data
|
|
if joined.count == count { completion(joined) }
|
|
else if complete { completion(nil) }
|
|
else { self?.receiveExactly(count, accumulated: joined, completion: completion) }
|
|
}
|
|
}
|
|
}
|