diff --git a/storage/blockexchange/engine/engine.nim b/storage/blockexchange/engine/engine.nim index ab97050a..0bb21190 100644 --- a/storage/blockexchange/engine/engine.nim +++ b/storage/blockexchange/engine/engine.nim @@ -1141,6 +1141,21 @@ proc new*( proc wantBlocksRequestHandler( peer: PeerId, req: WantBlocksRequest ): Future[seq[BlockDelivery]] {.async: (raises: [CancelledError]).} = + let maxIndex = high(Natural).uint64 + var totalCount: uint64 = 0 + for r in req.ranges: + if r.count == 0 or r.count > MaxBlocksPerBatch or r.start > maxIndex or + r.count - 1 > maxIndex - r.start or r.start > uint64.high - r.count or + r.count > uint64.high - totalCount: + warn "Rejecting WantBlocks request: invalid range", + peer = peer, start = r.start, count = r.count + return @[] + totalCount += r.count + if totalCount > MaxBlocksPerBatch: + warn "Rejecting WantBlocks request: total blocks exceeds cap", + peer = peer, total = totalCount + return @[] + var blockDeliveries: seq[BlockDelivery] notFoundCount = 0 diff --git a/tests/storage/blockexchange/engine/testengine.nim b/tests/storage/blockexchange/engine/testengine.nim index d8a18762..b806e52b 100644 --- a/tests/storage/blockexchange/engine/testengine.nim +++ b/tests/storage/blockexchange/engine/testengine.nim @@ -14,6 +14,7 @@ import pkg/storage/merkletree import pkg/storage/blockexchange/utils import pkg/storage/blockexchange/engine/activedownload {.all.} import pkg/storage/blockexchange/engine/downloadmanager {.all.} +import pkg/storage/blockexchange/protocol/constants import ../../../asynctest import ../../helpers @@ -268,6 +269,66 @@ asyncchecksuite "NetworkStore engine handlers": await engine.wantListHandler(peerId, wantList) await done + test "WantBlocks: rejects range with count = 0": + let req = + WantBlocksRequest(requestId: 1, treeCid: Cid.example, ranges: @[(0'u64, 0'u64)]) + + let blocks = await network.handlers.onWantBlocksRequest(peerId, req) + check blocks.len == 0 + + test "WantBlocks: rejects range with count > MaxBlocksPerBatch": + let req = WantBlocksRequest( + requestId: 1, + treeCid: Cid.example, + ranges: @[(0'u64, MaxBlocksPerBatch.uint64 + 1)], + ) + + let blocks = await network.handlers.onWantBlocksRequest(peerId, req) + check blocks.len == 0 + + test "WantBlocks: rejects range whose start+count overflows": + let req = WantBlocksRequest( + requestId: 1, treeCid: Cid.example, ranges: @[(uint64.high, 1'u64)] + ) + + let blocks = await network.handlers.onWantBlocksRequest(peerId, req) + check blocks.len == 0 + + test "WantBlocks: rejects range whose max index exceeds Natural": + let req = WantBlocksRequest( + requestId: 1, treeCid: Cid.example, ranges: @[(high(Natural).uint64 + 1, 1'u64)] + ) + + let blocks = await network.handlers.onWantBlocksRequest(peerId, req) + check blocks.len == 0 + + test "WantBlocks: rejects when total count across ranges exceeds cap": + var ranges: seq[tuple[start: uint64, count: uint64]] = @[] + let halfMaxBlocksPerBatchPlusOne = (MaxBlocksPerBatch div 2).uint64 + 1 + ranges.add((0'u64, halfMaxBlocksPerBatchPlusOne)) + ranges.add((10_000'u64, halfMaxBlocksPerBatchPlusOne)) + + let req = WantBlocksRequest(requestId: 1, treeCid: Cid.example, ranges: ranges) + + let blocks = await network.handlers.onWantBlocksRequest(peerId, req) + check blocks.len == 0 + + test "WantBlocks: accepts a valid small request": + let + tree = StorageMerkleTree.init(blocks.mapIt(it.cid)).tryGet + rootCid = tree.rootCid.tryGet() + + for i, blk in blocks: + (await localStore.putBlock(blk)).tryGet() + (await localStore.putCidAndProof(rootCid, i, blk.cid, tree.getProof(i).tryGet())).tryGet() + + let req = WantBlocksRequest( + requestId: 1, treeCid: rootCid, ranges: @[(0'u64, blocks.len.uint64)] + ) + + let delivered = await network.handlers.onWantBlocksRequest(peerId, req) + check delivered.len == blocks.len + suite "IsIndexInRanges": test "Empty ranges returns false": let ranges: seq[(uint64, uint64)] = @[]