smartban.go

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