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