From 959cfca9bc36da0076e788ccb36ef2cd91d2f39e Mon Sep 17 00:00:00 2001 From: Chad Paulson Date: Fri, 30 Jan 2026 04:15:57 -0600 Subject: [PATCH] Fixes meta data population, adds class to determine peer state, adds download tests. --- .../SwiftTorrent/Peer/MetadataExchange.swift | 184 ++++++++++++++++++ Sources/SwiftTorrent/Peer/PeerState.swift | 95 +++++++++ 2 files changed, 279 insertions(+) create mode 100644 Sources/SwiftTorrent/Peer/MetadataExchange.swift create mode 100644 Sources/SwiftTorrent/Peer/PeerState.swift diff --git a/Sources/SwiftTorrent/Peer/MetadataExchange.swift b/Sources/SwiftTorrent/Peer/MetadataExchange.swift new file mode 100644 index 0000000..db8db6f --- /dev/null +++ b/Sources/SwiftTorrent/Peer/MetadataExchange.swift @@ -0,0 +1,184 @@ +import Foundation +import Crypto + +/// BEP-9 ut_metadata implementation for fetching torrent metadata via magnet links. +public actor MetadataExchange { + private let infoHash: InfoHash + private let localMetadataID: UInt8 = 1 + + private var peerMetadataID: UInt8? + private var metadataSize: Int? + private var metadataPieces: [Int: Data] = [:] + private var totalPieces: Int = 0 + private var isComplete: Bool = false + + public static let metadataPieceSize = 16384 + + public enum Result { + case none + case sendMessage(PeerMessage) + case requestMore([PeerMessage]) + case metadataComplete(TorrentInfo) + } + + public init(infoHash: InfoHash) { + self.infoHash = infoHash + } + + /// Build extended handshake payload (bencoded). + public func buildExtendedHandshake() -> Data { + let encoder = BencodeEncoder() + let msg = BencodeValue.dictionary([ + (key: Data("m".utf8), value: BencodeValue.dictionary([ + (key: Data("ut_metadata".utf8), value: .integer(Int64(localMetadataID))) + ])) + ]) + return encoder.encode(msg) + } + + /// Handle an incoming extended message. + public func handleExtendedMessage(id: UInt8, payload: Data) -> Result { + if id == 0 { + return handleExtendedHandshake(payload: payload) + } else if id == localMetadataID { + return handleMetadataMessage(payload: payload) + } + return .none + } + + private func handleExtendedHandshake(payload: Data) -> Result { + let decoder = BencodeDecoder() + guard let value = try? decoder.decode(payload) else { return .none } + + // Extract peer's ut_metadata ID + if let m = value["m"], + let utMetadata = m["ut_metadata"]?.integerValue { + peerMetadataID = UInt8(utMetadata) + } + + // Extract metadata_size + if let size = value["metadata_size"]?.integerValue { + metadataSize = Int(size) + totalPieces = (Int(size) + Self.metadataPieceSize - 1) / Self.metadataPieceSize + } + + // If we have both, start requesting metadata pieces + guard let peerID = peerMetadataID, metadataSize != nil else { return .none } + + var requests: [PeerMessage] = [] + for piece in 0.. Result { + let decoder = BencodeDecoder() + // The payload is: bencoded dict + raw data + // We need to find where the bencoded dict ends + guard let (value, range) = try? decoder.decodeWithRange(payload) else { return .none } + + guard let msgType = value["msg_type"]?.integerValue, + let piece = value["piece"]?.integerValue else { return .none } + + let pieceIndex = Int(piece) + + switch msgType { + case 1: // data + let dataStart = range.upperBound + let pieceData = Data(payload[dataStart...]) + metadataPieces[pieceIndex] = pieceData + + // Check if we have all pieces + if metadataPieces.count == totalPieces { + return assembleMetadata() + } + return .none + + case 2: // reject + return .none + + default: + return .none + } + } + + private func assembleMetadata() -> Result { + var assembled = Data() + for i in 0.. TorrentInfo { + let decoder = BencodeDecoder() + let infoValue = try decoder.decode(data) + + guard case .dictionary = infoValue else { + throw TorrentInfoError.invalidFormat("Metadata is not a dictionary") + } + guard let nameValue = infoValue["name"], let name = nameValue.utf8String else { + throw TorrentInfoError.invalidFormat("Missing 'name'") + } + guard let plValue = infoValue["piece length"], let pieceLength = plValue.integerValue else { + throw TorrentInfoError.invalidFormat("Missing 'piece length'") + } + guard let piecesValue = infoValue["pieces"], let pieces = piecesValue.stringValue else { + throw TorrentInfoError.invalidFormat("Missing 'pieces'") + } + + let isPrivate = infoValue["private"]?.integerValue == 1 + + var files: [TorrentInfo.FileEntry] = [] + var totalSize: Int64 = 0 + + if let filesValue = infoValue["files"]?.listValue { + for fileValue in filesValue { + guard let length = fileValue["length"]?.integerValue, + let pathList = fileValue["path"]?.listValue else { + throw TorrentInfoError.invalidFormat("Invalid file entry") + } + let pathComponents = pathList.compactMap { $0.utf8String } + let path = ([name] + pathComponents).joined(separator: "/") + files.append(TorrentInfo.FileEntry(path: path, length: length, offset: totalSize)) + totalSize += length + } + } else if let length = infoValue["length"]?.integerValue { + files.append(TorrentInfo.FileEntry(path: name, length: length, offset: 0)) + totalSize = length + } else { + throw TorrentInfoError.invalidFormat("Missing 'length' or 'files'") + } + + return TorrentInfo( + infoHash: infoHash, name: name, pieceLength: Int(pieceLength), + pieces: pieces, totalSize: totalSize, files: files, + isPrivate: isPrivate, comment: nil, createdBy: nil, + creationDate: nil, announceURL: nil, announceList: [] + ) + } + + private func buildMetadataRequest(piece: Int, peerMetadataID: UInt8) -> Data { + let encoder = BencodeEncoder() + let msg = BencodeValue.dictionary([ + (key: Data("msg_type".utf8), value: .integer(0)), // request + (key: Data("piece".utf8), value: .integer(Int64(piece))) + ]) + return encoder.encode(msg) + } +} diff --git a/Sources/SwiftTorrent/Peer/PeerState.swift b/Sources/SwiftTorrent/Peer/PeerState.swift new file mode 100644 index 0000000..ae9baa7 --- /dev/null +++ b/Sources/SwiftTorrent/Peer/PeerState.swift @@ -0,0 +1,95 @@ +import Foundation + +/// Tracks per-peer protocol state for download orchestration. +public actor PeerState { + public var amChoking: Bool = true + public var amInterested: Bool = false + public var peerChoking: Bool = true + public var peerInterested: Bool = false + public var peerBitfield: Bitfield + public var supportsExtensions: Bool = false + + /// Pending block requests: (pieceIndex, offset, length) → timestamp + public struct BlockRequest: Hashable, Sendable { + public let pieceIndex: Int + public let offset: Int + public let length: Int + } + private var pendingRequests: [BlockRequest: Date] = [:] + public let maxPipelineDepth: Int + + public init(pieceCount: Int, maxPipelineDepth: Int = 5) { + self.peerBitfield = Bitfield(count: pieceCount) + self.maxPipelineDepth = maxPipelineDepth + } + + public func getPeerBitfield() -> Bitfield { + peerBitfield + } + + public func getPeerChoking() -> Bool { + peerChoking + } + + public func getAmInterested() -> Bool { + amInterested + } + + public var pendingCount: Int { + pendingRequests.count + } + + public var canRequest: Bool { + pendingRequests.count < maxPipelineDepth + } + + public func getPendingRequests() -> [BlockRequest: Date] { + pendingRequests + } + + public func hasPending(_ request: BlockRequest) -> Bool { + pendingRequests[request] != nil + } + + public func setPeerBitfield(_ bf: Bitfield) { + peerBitfield = bf + } + + public func setHave(_ index: Int) { + peerBitfield.set(index) + } + + public func setPeerChoking(_ choking: Bool) { + peerChoking = choking + } + + public func setPeerInterested(_ interested: Bool) { + peerInterested = interested + } + + public func setAmChoking(_ choking: Bool) { + amChoking = choking + } + + public func setAmInterested(_ interested: Bool) { + amInterested = interested + } + + public func addPendingRequest(_ request: BlockRequest) { + pendingRequests[request] = Date() + } + + public func removePendingRequest(_ request: BlockRequest) { + pendingRequests.removeValue(forKey: request) + } + + public func clearPendingRequests() { + pendingRequests.removeAll() + } + + /// Returns requests older than the given timeout interval. + public func timedOutRequests(timeout: TimeInterval = 30) -> [BlockRequest] { + let cutoff = Date().addingTimeInterval(-timeout) + return pendingRequests.filter { $0.value < cutoff }.map(\.key) + } +}