smartban/smartban.go

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