From bb404a2d13051dc258132d1b235fb4c11e16c907 Mon Sep 17 00:00:00 2001 From: Chad Paulson Date: Fri, 30 Jan 2026 04:16:12 -0600 Subject: [PATCH] Fixes meta data population, adds class to determine peer state, adds download tests. --- Tests/SwiftTorrentTests/DownloadTests.swift | 369 ++++++++++++++++++++ 1 file changed, 369 insertions(+) create mode 100644 Tests/SwiftTorrentTests/DownloadTests.swift diff --git a/Tests/SwiftTorrentTests/DownloadTests.swift b/Tests/SwiftTorrentTests/DownloadTests.swift new file mode 100644 index 0000000..4228223 --- /dev/null +++ b/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: [] + ) +}