smartban.go

   1 package torrent
   2
   3 import (
   4         "bytes"
   5         "crypto/sha1"
   6         "net/netip"
   7
   8         g "github.com/anacrolix/generics"
   9
  10         "github.com/anacrolix/torrent/smartban"
  11 )
  12
  13 type bannableAddr = netip.Addr
  14
  15 type smartBanCache = smartban.Cache[bannableAddr, RequestIndex, [sha1.Size]byte]
  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         blockBuffer bytes.Buffer
  23         chunkSize   int
  24 }
  25
  26 func (me *blockCheckingWriter) checkBlock() {
  27         b := me.blockBuffer.Next(me.chunkSize)
  28         for _, peer := range me.cache.CheckBlock(me.requestIndex, b) {
  29                 g.MakeMapIfNilAndSet(&me.badPeers, peer, struct{}{})
  30         }
  31         me.requestIndex++
  32 }
  33
  34 func (me *blockCheckingWriter) checkFullBlocks() {
  35         for me.blockBuffer.Len() >= me.chunkSize {
  36                 me.checkBlock()
  37         }
  38 }
  39
  40 func (me *blockCheckingWriter) Write(b []byte) (int, error) {
  41         n, err := me.blockBuffer.Write(b)
  42         if err != nil {
  43                 // bytes.Buffer.Write should never fail.
  44                 panic(err)
  45         }
  46         me.checkFullBlocks()
  47         return n, err
  48 }
  49
  50 // Check any remaining block data. Terminal pieces or piece sizes that don't divide into the chunk
  51 // size cleanly may leave fragments that should be checked.
  52 func (me *blockCheckingWriter) Flush() {
  53         for me.blockBuffer.Len() != 0 {
  54                 me.checkBlock()
  55         }
  56 }