diff --git a/BitTorrent/File Manager/TorrentFileManager.swift b/BitTorrent/File Manager/TorrentFileManager.swift index 4f96764..bdba957 100644 --- a/BitTorrent/File Manager/TorrentFileManager.swift +++ b/BitTorrent/File Manager/TorrentFileManager.swift @@ -137,7 +137,7 @@ extension TorrentFileManager { try? bitfield.toData().write(to: fileURL) } - static func loadSavedProgressBitfield(infoHash: Data) -> BitField? { + static func loadSavedProgressBitfield(infoHash: Data, size: Int) -> BitField? { let fileName = sanitizedFileName(from: infoHash) let documentsPath = NSSearchPathForDirectoriesInDomains(.cachesDirectory, .userDomainMask, @@ -145,7 +145,7 @@ extension TorrentFileManager { let documentsUrl = URL(fileURLWithPath: documentsPath, isDirectory: true) let fileURL = documentsUrl.appendingPathComponent(fileName, isDirectory: false) if let data = try? Data(contentsOf: fileURL) { - return BitField(data: data) + return BitField(data: data, size: size) } return nil } diff --git a/BitTorrent/Models/BitField.swift b/BitTorrent/Models/BitField.swift index 030f288..6855830 100644 --- a/BitTorrent/Models/BitField.swift +++ b/BitTorrent/Models/BitField.swift @@ -20,11 +20,12 @@ public struct BitField: Equatable { self.value = Array(repeating: false, count: size) } - init(data: Data) { - self.init(size: data.count*8) + init(data: Data, size: Int) { + precondition(BitField.dataMatchesSize(size, data: data), "The number of bytes doesn't match the given size") + self.init(size: size) for byteIndex in 0 ..< data.count { let byte = data.correctingIndicies[byteIndex] - for i in 0 ..< 8 { + for i in 0 ..< 8 where (byteIndex*8 + i) < size { if isNthBitSet(byte, n: i) { set(at: byteIndex*8 + i) } @@ -32,6 +33,11 @@ public struct BitField: Equatable { } } + private static func dataMatchesSize(_ size: Int, data: Data) -> Bool { + let numberOfBytes = numberOfBytesForSize(size) + return data.count == numberOfBytes + } + fileprivate func isNthBitSet(_ byte: UInt8, n: Int) -> Bool { let mask: [UInt8] = [128, 64, 32, 16, 8, 4, 2, 1] let maskN: UInt8 = mask[n] @@ -51,12 +57,7 @@ public struct BitField: Equatable { } func toData() -> Data { - let numberOfBytes: Int - if (size % 8) == 0 { - numberOfBytes = size / 8 - } else { - numberOfBytes = (size / 8) + 1 - } + let numberOfBytes = BitField.numberOfBytesForSize(size) var bytes: [UInt8] = [] for i in 0.. Int { + if (size % 8) == 0 { + return size / 8 + } else { + return (size / 8) + 1 + } + } + public static func ==(_ lhs: BitField, _ rhs: BitField) -> Bool { return lhs.value == rhs.value } @@ -105,7 +114,7 @@ extension BitField: Collection { } public subscript(position: Index) -> (index: Int, isSet: Bool) { - precondition(indices.contains(position), "out of bounds") + precondition(position >= 0 && position < size, "out of bounds") return (index: position, isSet: value[position]) } diff --git a/BitTorrent/Models/BitFieldTests.swift b/BitTorrent/Models/BitFieldTests.swift index 5b1506d..9a7f364 100644 --- a/BitTorrent/Models/BitFieldTests.swift +++ b/BitTorrent/Models/BitFieldTests.swift @@ -76,7 +76,7 @@ class BitFieldTests: XCTestCase { example.set(at: 14) let data = example.toData() - let result = BitField(data: data) + let result = BitField(data: data, size: 16) XCTAssertEqual(result, example) } diff --git a/BitTorrent/Peer/Communicator/TorrentPeerCommunicator.swift b/BitTorrent/Peer/Communicator/TorrentPeerCommunicator.swift index c6730fc..4d4ff2e 100644 --- a/BitTorrent/Peer/Communicator/TorrentPeerCommunicator.swift +++ b/BitTorrent/Peer/Communicator/TorrentPeerCommunicator.swift @@ -19,7 +19,7 @@ protocol TorrentPeerCommunicatorDelegate: class { func peerBecameInterested(_ sender: TorrentPeerCommunicator) func peerBecameUninterested(_ sender: TorrentPeerCommunicator) func peer(_ sender: TorrentPeerCommunicator, hasPiece piece: Int) - func peer(_ sender: TorrentPeerCommunicator, hasBitField bitField: BitField) + func peer(_ sender: TorrentPeerCommunicator, hasBitFieldData bitFieldData: Data) func peer(_ sender: TorrentPeerCommunicator, requestedPiece index: Int, begin: Int, length: Int) func peer(_ sender: TorrentPeerCommunicator, sentPiece index: Int, begin: Int, block: Data) func peer(_ sender: TorrentPeerCommunicator, cancelledRequestedPiece index: Int, begin: Int, length: Int) @@ -335,8 +335,7 @@ extension TorrentPeerCommunicator: TorrentPeerMessageBufferDelegate { private func processBitFieldMessage(_ message: Data) { let bitFieldData = message.correctingIndicies[5 ..< message.count] - let bitField = BitField(data: bitFieldData) - delegate?.peer(self, hasBitField: bitField) + delegate?.peer(self, hasBitFieldData: bitFieldData) } private func processRequestMessage(_ message: Data) { diff --git a/BitTorrent/Peer/Communicator/TorrentPeerCommunicatorDelegateStub.swift b/BitTorrent/Peer/Communicator/TorrentPeerCommunicatorDelegateStub.swift index e95585d..20ac293 100644 --- a/BitTorrent/Peer/Communicator/TorrentPeerCommunicatorDelegateStub.swift +++ b/BitTorrent/Peer/Communicator/TorrentPeerCommunicatorDelegateStub.swift @@ -79,10 +79,10 @@ class TorrentPeerCommunicatorDelegateStub: TorrentPeerCommunicatorDelegate { } var peerHasBitFieldCalled = false - var peerHasBitFieldParameters: (sender: TorrentPeerCommunicator, bitField: BitField)? - func peer(_ sender: TorrentPeerCommunicator, hasBitField bitField: BitField) { + var peerHasBitFieldParameters: (sender: TorrentPeerCommunicator, bitFieldData: Data)? + func peer(_ sender: TorrentPeerCommunicator, hasBitFieldData bitFieldData: Data) { peerHasBitFieldCalled = true - peerHasBitFieldParameters = (sender, bitField) + peerHasBitFieldParameters = (sender, bitFieldData) } var peerRequestedPieceCalled = false diff --git a/BitTorrent/Peer/Communicator/TorrentPeerCommunicatorReadTests.swift b/BitTorrent/Peer/Communicator/TorrentPeerCommunicatorReadTests.swift index 5d6952a..ac21abf 100644 --- a/BitTorrent/Peer/Communicator/TorrentPeerCommunicatorReadTests.swift +++ b/BitTorrent/Peer/Communicator/TorrentPeerCommunicatorReadTests.swift @@ -107,7 +107,7 @@ extension TorrentPeerComminicatorTests { XCTAssert(delegate.peerHasBitFieldCalled) XCTAssertEqual(delegate.peerHasBitFieldParameters?.sender, sut) - XCTAssertEqual(delegate.peerHasBitFieldParameters?.bitField, bitField) + XCTAssertEqual(delegate.peerHasBitFieldParameters?.bitFieldData, bitField.toData()) } func test_delegateCalled_whenPeerSendsRequest() { diff --git a/BitTorrent/Peer/TorrentPeer.swift b/BitTorrent/Peer/TorrentPeer.swift index 042b5ca..938b824 100644 --- a/BitTorrent/Peer/TorrentPeer.swift +++ b/BitTorrent/Peer/TorrentPeer.swift @@ -269,12 +269,14 @@ extension TorrentPeer: TorrentPeerCommunicatorDelegate { } func peer(_ sender: TorrentPeerCommunicator, hasPiece piece: Int) { + guard !currentProgress.isSet(at: piece) else { return } currentProgress.set(at: piece) delegate?.peerHasNewAvailablePieces(self) } - func peer(_ sender: TorrentPeerCommunicator, hasBitField bitField: BitField) { - currentProgress = bitField + func peer(_ sender: TorrentPeerCommunicator, hasBitFieldData bitFieldData: Data) { + let bitFieldSize = currentProgress.size + currentProgress = BitField(data: bitFieldData, size: bitFieldSize) delegate?.peerHasNewAvailablePieces(self) } diff --git a/BitTorrent/Peer/TorrentPeerDelegateStub.swift b/BitTorrent/Peer/TorrentPeerDelegateStub.swift index 0051413..e805558 100644 --- a/BitTorrent/Peer/TorrentPeerDelegateStub.swift +++ b/BitTorrent/Peer/TorrentPeerDelegateStub.swift @@ -10,10 +10,10 @@ class TorrentPeerDelegateStub: TorrentPeerDelegate { - var peerHasNewAvailablePiecesCalled = false + var peerHasNewAvailablePiecesCallCount = 0 var peerHasNewAvailablePiecesParameter: TorrentPeer? func peerHasNewAvailablePieces(_ sender: TorrentPeer) { - peerHasNewAvailablePiecesCalled = true + peerHasNewAvailablePiecesCallCount += 1 peerHasNewAvailablePiecesParameter = sender } diff --git a/BitTorrent/Peer/TorrentPeerTests.swift b/BitTorrent/Peer/TorrentPeerTests.swift index 701d467..8ae0dcd 100644 --- a/BitTorrent/Peer/TorrentPeerTests.swift +++ b/BitTorrent/Peer/TorrentPeerTests.swift @@ -111,25 +111,32 @@ class TorrentPeerTests: XCTestCase { } func test_delegateNotifiedAfterBitField() { - communicator.delegate?.peer(communicator, hasBitField: bitField) + communicator.delegate?.peer(communicator, hasBitFieldData: bitField.toData()) - XCTAssert(delegate.peerHasNewAvailablePiecesCalled) + XCTAssertEqual(delegate.peerHasNewAvailablePiecesCallCount, 1) XCTAssert(delegate.peerHasNewAvailablePiecesParameter === sut) } func test_delegateNotifiedAfterHaveMessage() { communicator.delegate?.peer(communicator, hasPiece: 0) - XCTAssert(delegate.peerHasNewAvailablePiecesCalled) + XCTAssertEqual(delegate.peerHasNewAvailablePiecesCallCount, 1) XCTAssert(delegate.peerHasNewAvailablePiecesParameter === sut) } + func test_delegateNotNotifiedAfterRedundantHaveMessage() { + communicator.delegate?.peer(communicator, hasPiece: 0) + communicator.delegate?.peer(communicator, hasPiece: 0) + + XCTAssertEqual(delegate.peerHasNewAvailablePiecesCallCount, 1) + } + // MARK: - Tracking peer status func test_bitFieldRecorded() { var bitField = BitField(size: 10) bitField.set(at: 0) - communicator.delegate?.peer(communicator, hasBitField: bitField) + communicator.delegate?.peer(communicator, hasBitFieldData: bitField.toData()) XCTAssertEqual(sut.currentProgress, bitField) } diff --git a/BitTorrent/Progress Manager/TorrentProgressManager.swift b/BitTorrent/Progress Manager/TorrentProgressManager.swift index 66b0e67..f211be6 100644 --- a/BitTorrent/Progress Manager/TorrentProgressManager.swift +++ b/BitTorrent/Progress Manager/TorrentProgressManager.swift @@ -23,11 +23,13 @@ class TorrentProgressManager { let downloadDirectory = rootDirectory + "/" + metaInfo.sensibleDownloadDirectoryName() let fileManager = TorrentFileManager(metaInfo: metaInfo, rootDirectory: downloadDirectory) + let bitFieldSize = metaInfo.info.pieces.count let progress: TorrentProgress - if let bitField = TorrentFileManager.loadSavedProgressBitfield(infoHash: metaInfo.infoHash) { + if let bitField = TorrentFileManager.loadSavedProgressBitfield(infoHash: metaInfo.infoHash, + size: bitFieldSize) { progress = TorrentProgress(bitField: bitField) } else { - progress = TorrentProgress(size: metaInfo.info.pieces.count) + progress = TorrentProgress(size: bitFieldSize) } self.init(fileManager: fileManager, progress: progress) }