smartban.go

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