]> Sergey Matveev's repositories - btrtrc.git/blob - smartban/smartban.go
Add check that torrent data cleanup works in TestIssue335
[btrtrc.git] / 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 }