Fixes meta data population, adds class to determine peer state, adds download tests.

2026-05-28 15:27:20 +00:00 · 2026-01-30 04:15:57 -06:00
parent 034053d3fb
commit 959cfca9bc
2 changed files with 279 additions and 0 deletions
@@ -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..<totalPieces {
+            let requestPayload = buildMetadataRequest(piece: piece, peerMetadataID: peerID)
+            requests.append(.extended(id: peerID, payload: requestPayload))
+        }
+        return .requestMore(requests)
+    }
+
+    private func handleMetadataMessage(payload: Data) -> 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..<totalPieces {
+            guard let piece = metadataPieces[i] else { return .none }
+            assembled.append(piece)
+        }
+
+        // Verify SHA-1 matches info hash
+        let hash = Data(Insecure.SHA1.hash(data: assembled))
+        guard hash == infoHash.bytes else {
+            metadataPieces.removeAll()
+            return .none
+        }
+
+        isComplete = true
+
+        // Parse into TorrentInfo
+        guard let info = try? parseInfoFromMetadata(assembled) else { return .none }
+        return .metadataComplete(info)
+    }
+
+    private func parseInfoFromMetadata(_ data: Data) throws -> 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)
+    }
+}
@@ -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)
+    }
+}