peer_protocol/handshake.go

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