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

Tests/SwiftTorrentTests/DownloadTests.swift

@@ -0,0 +1,369 @@
+import XCTest
+import Foundation
+import Crypto
+import NIOPosix
+@testable import SwiftTorrent
+
+final class PeerStateTests: XCTestCase {
+    func testInitialState() async {
+        let state = PeerState(pieceCount: 10)
+        let choking = await state.peerChoking
+        let interested = await state.peerInterested
+        let amChoking = await state.amChoking
+        let amInterested = await state.amInterested
+        XCTAssertTrue(choking)
+        XCTAssertFalse(interested)
+        XCTAssertTrue(amChoking)
+        XCTAssertFalse(amInterested)
+    }
+
+    func testChokeUnchokeTransitions() async {
+        let state = PeerState(pieceCount: 10)
+        await state.setPeerChoking(false)
+        var val = await state.getPeerChoking()
+        XCTAssertFalse(val)
+        await state.setPeerChoking(true)
+        val = await state.getPeerChoking()
+        XCTAssertTrue(val)
+    }
+
+    func testInterestedTransitions() async {
+        let state = PeerState(pieceCount: 10)
+        await state.setAmInterested(true)
+        var val = await state.getAmInterested()
+        XCTAssertTrue(val)
+        await state.setAmInterested(false)
+        val = await state.getAmInterested()
+        XCTAssertFalse(val)
+    }
+
+    func testPendingRequests() async {
+        let state = PeerState(pieceCount: 10, maxPipelineDepth: 3)
+        var canReq = await state.canRequest
+        XCTAssertTrue(canReq)
+        var count = await state.pendingCount
+        XCTAssertEqual(count, 0)
+
+        let r1 = PeerState.BlockRequest(pieceIndex: 0, offset: 0, length: 16384)
+        let r2 = PeerState.BlockRequest(pieceIndex: 0, offset: 16384, length: 16384)
+        let r3 = PeerState.BlockRequest(pieceIndex: 0, offset: 32768, length: 16384)
+
+        await state.addPendingRequest(r1)
+        await state.addPendingRequest(r2)
+        canReq = await state.canRequest
+        XCTAssertTrue(canReq)
+        count = await state.pendingCount
+        XCTAssertEqual(count, 2)
+
+        await state.addPendingRequest(r3)
+        canReq = await state.canRequest
+        XCTAssertFalse(canReq)
+
+        await state.removePendingRequest(r1)
+        canReq = await state.canRequest
+        XCTAssertTrue(canReq)
+    }
+
+    func testBitfield() async {
+        let state = PeerState(pieceCount: 10)
+        var bf = Bitfield(count: 10)
+        bf.set(3)
+        bf.set(7)
+        await state.setPeerBitfield(bf)
+
+        let peerBF = await state.getPeerBitfield()
+        XCTAssertTrue(peerBF.get(3))
+        XCTAssertTrue(peerBF.get(7))
+        XCTAssertFalse(peerBF.get(0))
+    }
+
+    func testSetHave() async {
+        let state = PeerState(pieceCount: 10)
+        await state.setHave(5)
+        let bf = await state.getPeerBitfield()
+        XCTAssertTrue(bf.get(5))
+        XCTAssertFalse(bf.get(4))
+    }
+
+    func testClearPendingRequests() async {
+        let state = PeerState(pieceCount: 10)
+        let r1 = PeerState.BlockRequest(pieceIndex: 0, offset: 0, length: 16384)
+        await state.addPendingRequest(r1)
+        var count = await state.pendingCount
+        XCTAssertEqual(count, 1)
+        await state.clearPendingRequests()
+        count = await state.pendingCount
+        XCTAssertEqual(count, 0)
+    }
+}
+
+final class BlockSplittingTests: XCTestCase {
+    func testStandardPieceBlockCount() async {
+        let info = makeTorrentInfo(pieceLength: 262144, totalSize: 1048576)
+        let pm = PieceManager(info: info)
+        let count = await pm.blockCount(for: 0)
+        XCTAssertEqual(count, 16)
+    }
+
+    func testLastPieceBlockCount() async {
+        let info = makeTorrentInfo(pieceLength: 262144, totalSize: 300000)
+        let pm = PieceManager(info: info)
+        let lastIdx = await pm.getPieceCount() - 1
+        let count = await pm.blockCount(for: lastIdx)
+        XCTAssertEqual(count, 3)
+    }
+
+    func testSmallPiece() async {
+        let info = makeTorrentInfo(pieceLength: 8192, totalSize: 8192)
+        let pm = PieceManager(info: info)
+        let count = await pm.blockCount(for: 0)
+        XCTAssertEqual(count, 1)
+    }
+}
+
+final class PieceManagerBlockTests: XCTestCase {
+    func testBlockReceiptAndCompletion() async {
+        let data = Data(repeating: 0xAB, count: 32768)
+        let hash = Data(Insecure.SHA1.hash(data: data))
+        let info = makeTorrentInfo(pieceLength: 32768, totalSize: 32768, pieceHashes: hash)
+        let pm = PieceManager(info: info)
+
+        await pm.startPiece(0)
+        let inProg1 = await pm.isInProgress(0)
+        XCTAssertTrue(inProg1)
+
+        await pm.addBlock(pieceIndex: 0, offset: 0, data: Data(repeating: 0xAB, count: 16384))
+        await pm.addBlock(pieceIndex: 0, offset: 16384, data: Data(repeating: 0xAB, count: 16384))
+
+        let buffer = await pm.getPieceBuffer(0)
+        XCTAssertEqual(buffer?.count, 32768)
+
+        let verified = await pm.completePiece(0)
+        XCTAssertTrue(verified)
+        let has = await pm.hasPiece(0)
+        XCTAssertTrue(has)
+        let inProg2 = await pm.isInProgress(0)
+        XCTAssertFalse(inProg2)
+    }
+
+    func testHashMismatchRejection() async {
+        let badHash = Data(repeating: 0x00, count: 20)
+        let info = makeTorrentInfo(pieceLength: 16384, totalSize: 16384, pieceHashes: badHash)
+        let pm = PieceManager(info: info)
+
+        await pm.startPiece(0)
+        await pm.addBlock(pieceIndex: 0, offset: 0, data: Data(repeating: 0xFF, count: 16384))
+
+        let verified = await pm.completePiece(0)
+        XCTAssertFalse(verified)
+        let has = await pm.hasPiece(0)
+        XCTAssertFalse(has)
+        let inProg = await pm.isInProgress(0)
+        XCTAssertFalse(inProg)
+    }
+}
+
+final class PiecePickerIntegrationTests: XCTestCase {
+    func testRarestFirstPicking() {
+        var picker = PiecePicker(pieceCount: 5)
+        var have = Bitfield(count: 5)
+
+        var peerA = Bitfield(count: 5)
+        peerA.set(0); peerA.set(1); peerA.set(2)
+        picker.addPeerBitfield(peerA)
+
+        var peerB = Bitfield(count: 5)
+        peerB.set(0); peerB.set(1)
+        picker.addPeerBitfield(peerB)
+
+        var peerC = Bitfield(count: 5)
+        peerC.set(0)
+        picker.addPeerBitfield(peerC)
+
+        let picked = picker.pick(have: have, peerHas: peerA)
+        XCTAssertEqual(picked, 2)
+
+        have.set(2)
+        let picked2 = picker.pick(have: have, peerHas: peerA)
+        XCTAssertEqual(picked2, 1)
+    }
+
+    func testPickMultiple() {
+        var picker = PiecePicker(pieceCount: 4)
+        var peerBF = Bitfield(count: 4)
+        peerBF.set(0); peerBF.set(1); peerBF.set(2); peerBF.set(3)
+        picker.addPeerBitfield(peerBF)
+
+        let have = Bitfield(count: 4)
+        let picks = picker.pickMultiple(have: have, peerHas: peerBF, count: 3)
+        XCTAssertEqual(picks.count, 3)
+    }
+}
+
+final class MetadataExchangeTests: XCTestCase {
+    func testExtendedHandshake() async {
+        let hash = InfoHash(bytes: Data(repeating: 0xAA, count: 20))
+        let metaEx = MetadataExchange(infoHash: hash)
+        let payload = await metaEx.buildExtendedHandshake()
+
+        let decoder = BencodeDecoder()
+        let value = try! decoder.decode(payload)
+        let utMeta = value["m"]?["ut_metadata"]?.integerValue
+        XCTAssertEqual(utMeta, 1)
+    }
+
+    func testMetadataAssemblyAndVerification() async {
+        let encoder = BencodeEncoder()
+        let infoDict = BencodeValue.dictionary([
+            (key: Data("length".utf8), value: .integer(1024)),
+            (key: Data("name".utf8), value: .string(Data("test.txt".utf8))),
+            (key: Data("piece length".utf8), value: .integer(512)),
+            (key: Data("pieces".utf8), value: .string(Data(repeating: 0, count: 40)))
+        ])
+        let infoData = encoder.encode(infoDict)
+        let infoHash = InfoHash.v1(from: infoData)
+
+        let metaEx = MetadataExchange(infoHash: infoHash)
+
+        let peerHandshake = encoder.encode(BencodeValue.dictionary([
+            (key: Data("m".utf8), value: BencodeValue.dictionary([
+                (key: Data("ut_metadata".utf8), value: .integer(2))
+            ])),
+            (key: Data("metadata_size".utf8), value: .integer(Int64(infoData.count)))
+        ]))
+
+        let result1 = await metaEx.handleExtendedMessage(id: 0, payload: peerHandshake)
+        if case .requestMore(let messages) = result1 {
+            XCTAssertEqual(messages.count, 1)
+        } else {
+            XCTFail("Expected requestMore")
+        }
+
+        let pieceResponse = encoder.encode(BencodeValue.dictionary([
+            (key: Data("msg_type".utf8), value: .integer(1)),
+            (key: Data("piece".utf8), value: .integer(0)),
+            (key: Data("total_size".utf8), value: .integer(Int64(infoData.count)))
+        ]))
+        var fullPayload = pieceResponse
+        fullPayload.append(infoData)
+
+        let result2 = await metaEx.handleExtendedMessage(id: 1, payload: fullPayload)
+        if case .metadataComplete(let torrentInfo) = result2 {
+            XCTAssertEqual(torrentInfo.name, "test.txt")
+            XCTAssertEqual(torrentInfo.totalSize, 1024)
+        } else {
+            XCTFail("Expected metadataComplete, got \(result2)")
+        }
+    }
+}
+
+final class DiskIOMultiFileTests: XCTestCase {
+    func testMultiFileWriteRead() async throws {
+        let tmpDir = NSTemporaryDirectory() + "swifttorrent_test_\(UUID().uuidString)"
+        defer { try? FileManager.default.removeItem(atPath: tmpDir) }
+
+        let files: [TorrentInfo.FileEntry] = [
+            .init(path: "test/a.bin", length: 100, offset: 0),
+            .init(path: "test/b.bin", length: 100, offset: 100)
+        ]
+        let fs = FileStorage(files: files, pieceLength: 64, totalSize: 200)
+        let dio = DiskIO(basePath: tmpDir, fileStorage: fs)
+
+        try await dio.allocateFiles()
+
+        let piece0 = Data(repeating: 0xAA, count: 64)
+        try await dio.writePiece(index: 0, data: piece0)
+
+        let piece1 = Data(repeating: 0xBB, count: 64)
+        try await dio.writePiece(index: 1, data: piece1)
+
+        let read0 = try await dio.readPiece(index: 0)
+        XCTAssertEqual(read0, piece0)
+
+        let read1 = try await dio.readPiece(index: 1)
+        XCTAssertEqual(read1, piece1)
+
+        let piece3 = Data(repeating: 0xDD, count: 8)
+        try await dio.writePiece(index: 3, data: piece3)
+        let read3 = try await dio.readPiece(index: 3)
+        XCTAssertEqual(read3, piece3)
+    }
+}
+
+final class ExtendedMessageTests: XCTestCase {
+    func testEncodeDecodeExtended() throws {
+        let payload = Data("test payload".utf8)
+        let msg = PeerMessage.extended(id: 1, payload: payload)
+        let encoded = msg.encode()
+        let decoded = try PeerMessage.decode(from: Data(encoded.dropFirst(4)))
+        XCTAssertEqual(decoded, msg)
+    }
+
+    func testExtendedKeepAliveStillWorks() throws {
+        let msg = PeerMessage.keepAlive
+        let encoded = msg.encode()
+        XCTAssertEqual(encoded.count, 4)
+    }
+}
+
+final class HandshakeExtensionBitTests: XCTestCase {
+    func testDefaultReservedHasExtensionBit() {
+        let hs = Handshake(infoHash: Data(repeating: 0, count: 20), peerID: Data(repeating: 0, count: 20))
+        XCTAssertEqual(hs.reserved[5] & 0x10, 0x10)
+    }
+
+    func testCustomReservedPreserved() {
+        let custom = Data(count: 8)
+        let hs = Handshake(infoHash: Data(repeating: 0, count: 20), peerID: Data(repeating: 0, count: 20), reserved: custom)
+        XCTAssertEqual(hs.reserved[5] & 0x10, 0)
+    }
+}
+
+final class TorrentHandleGetFilesTests: XCTestCase {
+    func testGetFilesReturnsNilWithoutInfo() async throws {
+        let settings = SessionSettings(listenPort: 0, savePath: NSTemporaryDirectory())
+        let params = try AddTorrentParams.fromMagnet(
+            "magnet:?xt=urn:btih:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa&dn=Test",
+            savePath: NSTemporaryDirectory()
+        )
+        let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { try? group.syncShutdownGracefully() }
+        let handle = TorrentHandle(params: params, settings: settings, group: group)
+        let files = await handle.getFiles()
+        XCTAssertNil(files)
+    }
+
+    func testGetFilesReturnsFilesWithInfo() async throws {
+        let info = makeTorrentInfo(pieceLength: 16384, totalSize: 32768)
+        let params = AddTorrentParams(torrentInfo: info, savePath: NSTemporaryDirectory())
+        let settings = SessionSettings(listenPort: 0, savePath: NSTemporaryDirectory())
+        let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { try? group.syncShutdownGracefully() }
+        let handle = TorrentHandle(params: params, settings: settings, group: group)
+        let files = await handle.getFiles()
+        XCTAssertNotNil(files)
+        XCTAssertEqual(files?.count, 1)
+        XCTAssertEqual(files?.first?.path, "test")
+    }
+}
+
+// MARK: - Helpers
+
+func makeTorrentInfo(pieceLength: Int, totalSize: Int64, pieceHashes: Data? = nil) -> TorrentInfo {
+    let pieceCount = Int((totalSize + Int64(pieceLength) - 1) / Int64(pieceLength))
+    let hashes = pieceHashes ?? Data(repeating: 0, count: pieceCount * 20)
+    return TorrentInfo(
+        infoHash: InfoHash(bytes: Data(repeating: 0, count: 20)),
+        name: "test",
+        pieceLength: pieceLength,
+        pieces: hashes,
+        totalSize: totalSize,
+        files: [TorrentInfo.FileEntry(path: "test", length: totalSize, offset: 0)],
+        isPrivate: false,
+        comment: nil,
+        createdBy: nil,
+        creationDate: nil,
+        announceURL: nil,
+        announceList: []
+    )
+}