limit max webseed response size

internal/piecepicker/piecepicker.go

@@ -37,6 +37,7 @@ type PiecePicker struct {
  piecesByAvailability []*myPiece
  piecesByStalled []*myPiece
  maxDuplicateDownload int
+ maxWebseedPieces int
  available uint32
  endgame bool
 }
@@ -87,11 +88,16 @@ func New(pieces []piece.Piece, maxDuplicateDownload int, webseedSources []*webse
  sps[i] = &ps[i]
  sps2[i] = &ps[i]
  }
+ maxWebseedPieces := len(pieces) / 20 // Download 5% of pieces in a single HTTP request (see BEP19)
+ if maxWebseedPieces == 0 {
+ maxWebseedPieces = 1
+ }
  return &PiecePicker{
  pieces: ps,
  piecesByAvailability: sps,
  piecesByStalled: sps2,
  maxDuplicateDownload: maxDuplicateDownload,
+ maxWebseedPieces: maxWebseedPieces,
  webseedSources: webseedSources,
  }
 }

internal/piecepicker/webseed.go

@@ -110,21 +110,21 @@ func (p *PiecePicker) findGaps() []Range {
 
 func (p *PiecePicker) findGapsWithDuplicate(duplicate bool) []Range {
  a := make([]Range, 0, len(p.pieces)/2)
- var inGap bool
+ var inGap bool // See BEP19 for definition of "gap".
  var begin uint32
  for _, pi := range p.pieces {
  if !inGap {
  if pi.AvailableForWebseed(duplicate) {
  begin = pi.Index
  inGap = true
- } else {
- continue
  }
  } else {
- if pi.AvailableForWebseed(duplicate) {
- continue
- } else {
- a = append(a, Range{Begin: begin, End: pi.Index})
+ r := Range{Begin: begin, End: pi.Index}
+ if r.Len() == p.maxWebseedPieces {
+ a = append(a, r)
+ begin = pi.Index
+ } else if !pi.AvailableForWebseed(duplicate) {
+ a = append(a, r)
  inGap = false
  }
  }
@@ -164,6 +164,6 @@ type Range struct {
 }
 
 // Len returns the number of pieces in the range.
-func (r Range) Len() uint32 {
- return r.End - r.Begin
+func (r Range) Len() int {
+ return int(r.End) - int(r.Begin)
 }

internal/piecepicker/webseed_test.go

@@ -8,12 +8,47 @@ import (
 )
 
 func TestPickLastPieceOfSmallestGap(t *testing.T) {
- pieces := make([]piece.Piece, numPieces)
+ pieces := make([]piece.Piece, 10)
  for i := range pieces {
  pieces[i] = newPiece(i)
  }
  pieces[1].Done = true
  peer := newPeer(0)
  pp := New(pieces, 2, nil)
+ pp.maxWebseedPieces = 10
  assert.Nil(t, pp.pickLastPieceOfSmallestGap(peer))
 }
+
+func TestFindGapsContiguous(t *testing.T) {
+ pieces := make([]piece.Piece, 10)
+ for i := range pieces {
+ p := newPiece(i)
+ pieces[i] = p
+ }
+ pp := New(pieces, 2, nil)
+ pp.maxWebseedPieces = 10
+ assert.Equal(t, []Range{{0, 10}}, pp.findGaps())
+}
+
+func TestFindGapsSplit(t *testing.T) {
+ pieces := make([]piece.Piece, 9)
+ for i := range pieces {
+ p := newPiece(i)
+ pieces[i] = p
+ }
+ pieces[4].Done = true
+ pp := New(pieces, 2, nil)
+ pp.maxWebseedPieces = 10
+ assert.Equal(t, []Range{{0, 4}, {5, 9}}, pp.findGaps())
+}
+
+func TestFindGapsLimitMaxPieces(t *testing.T) {
+ pieces := make([]piece.Piece, 10)
+ for i := range pieces {
+ p := newPiece(i)
+ pieces[i] = p
+ }
+ pp := New(pieces, 2, nil)
+ pp.maxWebseedPieces = 4
+ assert.Equal(t, []Range{{0, 4}, {4, 8}, {8, 10}}, pp.findGaps())
+}