Ensure the bitfield size is correct

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
     }

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..<numberOfBytes {
@@ -87,6 +88,14 @@ public struct BitField: Equatable {
         return Data(bytes: bytes)
     }
     
+    private static func numberOfBytesForSize(_ size: Int) -> 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])
     }
     

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)
     }
     

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) {

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

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() {

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)
     }
     

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
     }
     

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)
     }
     

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)
     }