peer_protocol/handshake.go

   1 package peer_protocol
   2
   3 import (
   4         "encoding/hex"
   5         "fmt"
   6         "io"
   7
   8         "golang.org/x/xerrors"
   9
  10         "github.com/anacrolix/missinggo/v2"
  11         "github.com/anacrolix/torrent/metainfo"
  12 )
  13
  14 type ExtensionBit uint
  15
  16 const (
  17         ExtensionBitDHT      = 0  // http://www.bittorrent.org/beps/bep_0005.html
  18         ExtensionBitExtended = 20 // http://www.bittorrent.org/beps/bep_0010.html
  19         ExtensionBitFast     = 2  // http://www.bittorrent.org/beps/bep_0006.html
  20 )
  21
  22 func handshakeWriter(w io.Writer, bb <-chan []byte, done chan<- error) {
  23         var err error
  24         for b := range bb {
  25                 _, err = w.Write(b)
  26                 if err != nil {
  27                         break
  28                 }
  29         }
  30         done <- err
  31 }
  32
  33 type (
  34         PeerExtensionBits [8]byte
  35 )
  36
  37 func (me PeerExtensionBits) String() string {
  38         return hex.EncodeToString(me[:])
  39 }
  40
  41 func NewPeerExtensionBytes(bits ...ExtensionBit) (ret PeerExtensionBits) {
  42         for _, b := range bits {
  43                 ret.SetBit(b, true)
  44         }
  45         return
  46 }
  47
  48 func (pex PeerExtensionBits) SupportsExtended() bool {
  49         return pex.GetBit(ExtensionBitExtended)
  50 }
  51
  52 func (pex PeerExtensionBits) SupportsDHT() bool {
  53         return pex.GetBit(ExtensionBitDHT)
  54 }
  55
  56 func (pex PeerExtensionBits) SupportsFast() bool {
  57         return pex.GetBit(ExtensionBitFast)
  58 }
  59
  60 func (pex *PeerExtensionBits) SetBit(bit ExtensionBit, on bool) {
  61         if on {
  62                 pex[7-bit/8] |= 1 << (bit % 8)
  63         } else {
  64                 pex[7-bit/8] &^= 1 << (bit % 8)
  65         }
  66 }
  67
  68 func (pex PeerExtensionBits) GetBit(bit ExtensionBit) bool {
  69         return pex[7-bit/8]&(1<<(bit%8)) != 0
  70 }
  71
  72 type HandshakeResult struct {
  73         PeerExtensionBits
  74         PeerID [20]byte
  75         metainfo.Hash
  76 }
  77
  78 // ih is nil if we expect the peer to declare the InfoHash, such as when the peer initiated the
  79 // connection. Returns ok if the Handshake was successful, and err if there was an unexpected
  80 // condition other than the peer simply abandoning the Handshake.
  81 func Handshake(
  82         sock io.ReadWriter, ih *metainfo.Hash, peerID [20]byte, extensions PeerExtensionBits,
  83 ) (
  84         res HandshakeResult, err error,
  85 ) {
  86         // Bytes to be sent to the peer. Should never block the sender.
  87         postCh := make(chan []byte, 4)
  88         // A single error value sent when the writer completes.
  89         writeDone := make(chan error, 1)
  90         // Performs writes to the socket and ensures posts don't block.
  91         go handshakeWriter(sock, postCh, writeDone)
  92
  93         defer func() {
  94                 close(postCh) // Done writing.
  95                 if err != nil {
  96                         return
  97                 }
  98                 // Wait until writes complete before returning from handshake.
  99                 err = <-writeDone
 100                 if err != nil {
 101                         err = fmt.Errorf("error writing: %s", err)
 102                 }
 103         }()
 104
 105         post := func(bb []byte) {
 106                 select {
 107                 case postCh <- bb:
 108                 default:
 109                         panic("mustn't block while posting")
 110                 }
 111         }
 112
 113         post([]byte(Protocol))
 114         post(extensions[:])
 115         if ih != nil { // We already know what we want.
 116                 post(ih[:])
 117                 post(peerID[:])
 118         }
 119         var b [68]byte
 120         _, err = io.ReadFull(sock, b[:68])
 121         if err != nil {
 122                 err = xerrors.Errorf("while reading: %w", err)
 123                 return
 124         }
 125         if string(b[:20]) != Protocol {
 126                 err = xerrors.Errorf("unexpected protocol string")
 127                 return
 128         }
 129         missinggo.CopyExact(&res.PeerExtensionBits, b[20:28])
 130         missinggo.CopyExact(&res.Hash, b[28:48])
 131         missinggo.CopyExact(&res.PeerID, b[48:68])
 132         // peerExtensions.Add(res.PeerExtensionBits.String(), 1)
 133
 134         // TODO: Maybe we can just drop peers here if we're not interested. This
 135         // could prevent them trying to reconnect, falsely believing there was
 136         // just a problem.
 137         if ih == nil { // We were waiting for the peer to tell us what they wanted.
 138                 post(res.Hash[:])
 139                 post(peerID[:])
 140         }
 141
 142         return
 143 }