peer_protocol/protocol.go

   1 package peer_protocol
   2
   3 import (
   4         "bufio"
   5         "bytes"
   6         "encoding/binary"
   7         "errors"
   8         "fmt"
   9         "io"
  10         "io/ioutil"
  11 )
  12
  13 type (
  14         MessageType byte
  15         Integer     uint32
  16 )
  17
  18 func (i *Integer) Read(r io.Reader) error {
  19         return binary.Read(r, binary.BigEndian, i)
  20 }
  21
  22 const (
  23         Protocol = "\x13BitTorrent protocol"
  24 )
  25
  26 const (
  27         Choke         MessageType = iota
  28         Unchoke                   // 1
  29         Interested                // 2
  30         NotInterested             // 3
  31         Have                      // 4
  32         Bitfield                  // 5
  33         Request                   // 6
  34         Piece                     // 7
  35         Cancel                    // 8
  36         Port                      // 9
  37
  38         // BEP 6
  39         Suggest     = 0xd  // 13
  40         HaveAll     = 0xe  // 14
  41         HaveNone    = 0xf  // 15
  42         Reject      = 0x10 // 16
  43         AllowedFast = 0x11 // 17
  44
  45         Extended = 20
  46
  47         HandshakeExtendedID = 0
  48
  49         RequestMetadataExtensionMsgType = 0
  50         DataMetadataExtensionMsgType    = 1
  51         RejectMetadataExtensionMsgType  = 2
  52 )
  53
  54 type Message struct {
  55         Keepalive            bool
  56         Type                 MessageType
  57         Index, Begin, Length Integer
  58         Piece                []byte
  59         Bitfield             []bool
  60         ExtendedID           byte
  61         ExtendedPayload      []byte
  62         Port                 uint16
  63 }
  64
  65 func (msg Message) MarshalBinary() (data []byte, err error) {
  66         buf := &bytes.Buffer{}
  67         if !msg.Keepalive {
  68                 err = buf.WriteByte(byte(msg.Type))
  69                 if err != nil {
  70                         return
  71                 }
  72                 switch msg.Type {
  73                 case Choke, Unchoke, Interested, NotInterested, HaveAll, HaveNone:
  74                 case Have:
  75                         err = binary.Write(buf, binary.BigEndian, msg.Index)
  76                 case Request, Cancel, Reject:
  77                         for _, i := range []Integer{msg.Index, msg.Begin, msg.Length} {
  78                                 err = binary.Write(buf, binary.BigEndian, i)
  79                                 if err != nil {
  80                                         break
  81                                 }
  82                         }
  83                 case Bitfield:
  84                         _, err = buf.Write(marshalBitfield(msg.Bitfield))
  85                 case Piece:
  86                         for _, i := range []Integer{msg.Index, msg.Begin} {
  87                                 err = binary.Write(buf, binary.BigEndian, i)
  88                                 if err != nil {
  89                                         return
  90                                 }
  91                         }
  92                         n, err := buf.Write(msg.Piece)
  93                         if err != nil {
  94                                 break
  95                         }
  96                         if n != len(msg.Piece) {
  97                                 panic(n)
  98                         }
  99                 case Extended:
 100                         err = buf.WriteByte(msg.ExtendedID)
 101                         if err != nil {
 102                                 return
 103                         }
 104                         _, err = buf.Write(msg.ExtendedPayload)
 105                 case Port:
 106                         err = binary.Write(buf, binary.BigEndian, msg.Port)
 107                 default:
 108                         err = fmt.Errorf("unknown message type: %v", msg.Type)
 109                 }
 110         }
 111         data = make([]byte, 4+buf.Len())
 112         binary.BigEndian.PutUint32(data, uint32(buf.Len()))
 113         if buf.Len() != copy(data[4:], buf.Bytes()) {
 114                 panic("bad copy")
 115         }
 116         return
 117 }
 118
 119 type Decoder struct {
 120         R         *bufio.Reader
 121         MaxLength Integer // TODO: Should this include the length header or not?
 122 }
 123
 124 // io.EOF is returned if the source terminates cleanly on a message boundary.
 125 func (d *Decoder) Decode(msg *Message) (err error) {
 126         var length Integer
 127         err = binary.Read(d.R, binary.BigEndian, &length)
 128         if err != nil {
 129                 if err != io.EOF {
 130                         err = fmt.Errorf("error reading message length: %s", err)
 131                 }
 132                 return
 133         }
 134         if length > d.MaxLength {
 135                 return errors.New("message too long")
 136         }
 137         if length == 0 {
 138                 msg.Keepalive = true
 139                 return
 140         }
 141         msg.Keepalive = false
 142         b := make([]byte, length)
 143         _, err = io.ReadFull(d.R, b)
 144         if err != nil {
 145                 if err == io.EOF {
 146                         err = io.ErrUnexpectedEOF
 147                 }
 148                 if err != io.ErrUnexpectedEOF {
 149                         err = fmt.Errorf("error reading message: %s", err)
 150                 }
 151                 return
 152         }
 153         r := bytes.NewReader(b)
 154         // Check that all of r was utilized.
 155         defer func() {
 156                 if err != nil {
 157                         return
 158                 }
 159                 if r.Len() != 0 {
 160                         err = fmt.Errorf("%d bytes unused in message type %d", r.Len(), msg.Type)
 161                 }
 162         }()
 163         msg.Keepalive = false
 164         c, err := r.ReadByte()
 165         if err != nil {
 166                 return
 167         }
 168         msg.Type = MessageType(c)
 169         switch msg.Type {
 170         case Choke, Unchoke, Interested, NotInterested, HaveAll, HaveNone:
 171                 return
 172         case Have:
 173                 err = msg.Index.Read(r)
 174         case Request, Cancel, Reject:
 175                 for _, data := range []*Integer{&msg.Index, &msg.Begin, &msg.Length} {
 176                         err = data.Read(r)
 177                         if err != nil {
 178                                 break
 179                         }
 180                 }
 181         case Bitfield:
 182                 b := make([]byte, length-1)
 183                 _, err = io.ReadFull(r, b)
 184                 msg.Bitfield = unmarshalBitfield(b)
 185         case Piece:
 186                 for _, pi := range []*Integer{&msg.Index, &msg.Begin} {
 187                         err = pi.Read(r)
 188                         if err != nil {
 189                                 break
 190                         }
 191                 }
 192                 if err != nil {
 193                         break
 194                 }
 195                 msg.Piece, err = ioutil.ReadAll(r)
 196         case Extended:
 197                 msg.ExtendedID, err = r.ReadByte()
 198                 if err != nil {
 199                         break
 200                 }
 201                 msg.ExtendedPayload, err = ioutil.ReadAll(r)
 202         case Port:
 203                 err = binary.Read(r, binary.BigEndian, &msg.Port)
 204         default:
 205                 err = fmt.Errorf("unknown message type %#v", c)
 206         }
 207         return
 208 }
 209
 210 type Bytes []byte
 211
 212 func (b Bytes) MarshalBinary() ([]byte, error) {
 213         return b, nil
 214 }
 215
 216 func unmarshalBitfield(b []byte) (bf []bool) {
 217         for _, c := range b {
 218                 for i := 7; i >= 0; i-- {
 219                         bf = append(bf, (c>>uint(i))&1 == 1)
 220                 }
 221         }
 222         return
 223 }
 224
 225 func marshalBitfield(bf []bool) (b []byte) {
 226         b = make([]byte, (len(bf)+7)/8)
 227         for i, have := range bf {
 228                 if !have {
 229                         continue
 230                 }
 231                 c := b[i/8]
 232                 c |= 1 << uint(7-i%8)
 233                 b[i/8] = c
 234         }
 235         return
 236 }