smartban/smartban.go

   1 package smartban
   2
   3 import (
   4         "sync"
   5
   6         g "github.com/anacrolix/generics"
   7 )
   8
   9 type Cache[Peer, BlockKey, Hash comparable] struct {
  10         Hash func([]byte) Hash
  11
  12         lock   sync.RWMutex
  13         blocks map[BlockKey][]peerAndHash[Peer, Hash]
  14 }
  15
  16 type Block[Key any] struct {
  17         Key  Key
  18         Data []byte
  19 }
  20
  21 type peerAndHash[Peer, Hash any] struct {
  22         Peer Peer
  23         Hash Hash
  24 }
  25
  26 func (me *Cache[Peer, BlockKey, Hash]) Init() {
  27         g.MakeMap(&me.blocks)
  28 }
  29
  30 func (me *Cache[Peer, BlockKey, Hash]) RecordBlock(peer Peer, key BlockKey, data []byte) {
  31         hash := me.Hash(data)
  32         me.lock.Lock()
  33         defer me.lock.Unlock()
  34         peers := me.blocks[key]
  35         peers = append(peers, peerAndHash[Peer, Hash]{peer, hash})
  36         me.blocks[key] = peers
  37 }
  38
  39 func (me *Cache[Peer, BlockKey, Hash]) CheckBlock(key BlockKey, data []byte) (bad []Peer) {
  40         correct := me.Hash(data)
  41         me.lock.RLock()
  42         defer me.lock.RUnlock()
  43         for _, item := range me.blocks[key] {
  44                 if item.Hash != correct {
  45                         bad = append(bad, item.Peer)
  46                 }
  47         }
  48         return
  49 }
  50
  51 func (me *Cache[Peer, BlockKey, Hash]) ForgetBlock(key BlockKey) {
  52         me.lock.Lock()
  53         defer me.lock.Unlock()
  54         delete(me.blocks, key)
  55 }
  56
  57 func (me *Cache[Peer, BlockKey, Hash]) HasBlocks() bool {
  58         me.lock.RLock()
  59         defer me.lock.RUnlock()
  60         return len(me.blocks) != 0
  61 }