github.com/anacrolix/torrent@v1.61.0/smartban.go (about)

     1  package torrent
     2  
     3  import (
     4  	"net/netip"
     5  
     6  	g "github.com/anacrolix/generics"
     7  	"github.com/anacrolix/missinggo/v2/panicif"
     8  
     9  	"github.com/anacrolix/torrent/smartban"
    10  )
    11  
    12  type bannableAddr = netip.Addr
    13  
    14  // TODO: Should be keyed on weak[Peer].
    15  type smartBanCache = smartban.Cache[bannableAddr, RequestIndex, uint64]
    16  
    17  type blockCheckingWriter struct {
    18  	cache        *smartBanCache
    19  	requestIndex RequestIndex
    20  	// Peers that didn't match blocks written now.
    21  	badPeers    map[bannableAddr]struct{}
    22  	chunkBuffer []byte
    23  	bufferUsed  int
    24  }
    25  
    26  func (me *blockCheckingWriter) chunkSize() int {
    27  	return len(me.chunkBuffer)
    28  }
    29  
    30  func (me *blockCheckingWriter) checkBufferedBlock() {
    31  	panicif.Zero(me.bufferUsed)
    32  	b := me.chunkBuffer[:me.bufferUsed]
    33  	me.bufferUsed = 0
    34  	me.checkBlock(b)
    35  }
    36  
    37  func (me *blockCheckingWriter) checkBlock(b []byte) {
    38  	for _, peer := range me.cache.CheckBlock(me.requestIndex, b) {
    39  		g.MakeMapIfNil(&me.badPeers)
    40  		me.badPeers[peer] = struct{}{}
    41  	}
    42  	me.requestIndex++
    43  }
    44  
    45  func (me *blockCheckingWriter) finishPartialBlock(b []byte) int {
    46  	if me.bufferUsed == 0 {
    47  		return 0
    48  	}
    49  	panicif.NotEq(len(me.chunkBuffer), me.chunkSize())
    50  	n := copy(me.chunkBuffer[me.bufferUsed:], b)
    51  	me.bufferUsed += n
    52  	if me.bufferUsed >= me.chunkSize() {
    53  		me.checkBufferedBlock()
    54  	}
    55  	return n
    56  }
    57  
    58  func (me *blockCheckingWriter) Write(b []byte) (n int, err error) {
    59  	n = me.finishPartialBlock(b)
    60  	b = b[n:]
    61  	if len(b) == 0 {
    62  		return
    63  	}
    64  	for len(b) >= me.chunkSize() {
    65  		me.checkBlock(b[:me.chunkSize()])
    66  		b = b[me.chunkSize():]
    67  		n += me.chunkSize()
    68  	}
    69  	panicif.NotZero(me.bufferUsed)
    70  	me.bufferUsed = copy(me.chunkBuffer, b)
    71  	n += me.bufferUsed
    72  	return n, err
    73  }
    74  
    75  // Check any remaining block data. Terminal pieces or piece sizes that don't divide into the chunk
    76  // size cleanly may leave fragments that should be checked.
    77  func (me *blockCheckingWriter) Flush() {
    78  	if me.bufferUsed != 0 {
    79  		me.checkBufferedBlock()
    80  	}
    81  	panicif.NotZero(me.bufferUsed)
    82  }