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

     1  package smartban
     2  
     3  import (
     4  	"iter"
     5  	"sync"
     6  
     7  	g "github.com/anacrolix/generics"
     8  )
     9  
    10  type Cache[Peer, BlockKey, Hash comparable] struct {
    11  	Hash func([]byte) Hash
    12  
    13  	// Wonder if we should make this an atomic.
    14  	lock   sync.RWMutex
    15  	blocks map[BlockKey][]peerAndHash[Peer, Hash]
    16  }
    17  
    18  type Block[Key any] struct {
    19  	Key  Key
    20  	Data []byte
    21  }
    22  
    23  type peerAndHash[Peer, Hash any] struct {
    24  	Peer Peer
    25  	Hash Hash
    26  }
    27  
    28  func (me *Cache[Peer, BlockKey, Hash]) Init() {
    29  	g.MakeMap(&me.blocks)
    30  }
    31  
    32  func (me *Cache[Peer, BlockKey, Hash]) RecordBlock(peer Peer, key BlockKey, data []byte) {
    33  	hash := me.Hash(data)
    34  	me.lock.Lock()
    35  	defer me.lock.Unlock()
    36  	peers := me.blocks[key]
    37  	peers = append(peers, peerAndHash[Peer, Hash]{peer, hash})
    38  	me.blocks[key] = peers
    39  }
    40  
    41  func (me *Cache[Peer, BlockKey, Hash]) CheckBlock(key BlockKey, data []byte) (bad []Peer) {
    42  	correct := me.Hash(data)
    43  	me.lock.RLock()
    44  	defer me.lock.RUnlock()
    45  	for _, item := range me.blocks[key] {
    46  		if item.Hash != correct {
    47  			bad = append(bad, item.Peer)
    48  		}
    49  	}
    50  	return
    51  }
    52  
    53  func (me *Cache[Peer, BlockKey, Hash]) ForgetBlockSeq(seq iter.Seq[BlockKey]) {
    54  	me.lock.Lock()
    55  	defer me.lock.Unlock()
    56  	if len(me.blocks) == 0 {
    57  		return
    58  	}
    59  	for key := range seq {
    60  		delete(me.blocks, key)
    61  	}
    62  }
    63  
    64  // Returns whether any block in the sequence has at least one peer recorded.
    65  func (me *Cache[Peer, BlockKey, Hash]) HasPeerForBlocks(seq iter.Seq[BlockKey]) bool {
    66  	me.lock.RLock()
    67  	defer me.lock.RUnlock()
    68  	if len(me.blocks) == 0 {
    69  		return false
    70  	}
    71  	for key := range seq {
    72  		if len(me.blocks[key]) != 0 {
    73  			return true
    74  		}
    75  	}
    76  	return false
    77  }
    78  
    79  func (me *Cache[Peer, BlockKey, Hash]) HasBlocks() bool {
    80  	me.lock.RLock()
    81  	defer me.lock.RUnlock()
    82  	return len(me.blocks) != 0
    83  }