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         Extended      = 20
  37
  38         HandshakeExtendedID = 0
  39
  40         RequestMetadataExtensionMsgType = 0
  41         DataMetadataExtensionMsgType    = 1
  42         RejectMetadataExtensionMsgType  = 2
  43 )
  44
  45 type Message struct {
  46         Keepalive            bool
  47         Type                 MessageType
  48         Index, Begin, Length Integer
  49         Piece                []byte
  50         Bitfield             []bool
  51         ExtendedID           byte
  52         ExtendedPayload      []byte
  53 }
  54
  55 func (msg Message) MarshalBinary() (data []byte, err error) {
  56         buf := &bytes.Buffer{}
  57         if !msg.Keepalive {
  58                 err = buf.WriteByte(byte(msg.Type))
  59                 if err != nil {
  60                         return
  61                 }
  62                 switch msg.Type {
  63                 case Choke, Unchoke, Interested, NotInterested:
  64                 case Have:
  65                         err = binary.Write(buf, binary.BigEndian, msg.Index)
  66                 case Request, Cancel:
  67                         for _, i := range []Integer{msg.Index, msg.Begin, msg.Length} {
  68                                 err = binary.Write(buf, binary.BigEndian, i)
  69                                 if err != nil {
  70                                         break
  71                                 }
  72                         }
  73                 case Bitfield:
  74                         _, err = buf.Write(marshalBitfield(msg.Bitfield))
  75                 case Piece:
  76                         for _, i := range []Integer{msg.Index, msg.Begin} {
  77                                 err = binary.Write(buf, binary.BigEndian, i)
  78                                 if err != nil {
  79                                         return
  80                                 }
  81                         }
  82                         n, err := buf.Write(msg.Piece)
  83                         if err != nil {
  84                                 break
  85                         }
  86                         if n != len(msg.Piece) {
  87                                 panic(n)
  88                         }
  89                 case Extended:
  90                         err = buf.WriteByte(msg.ExtendedID)
  91                         if err != nil {
  92                                 return
  93                         }
  94                         _, err = buf.Write(msg.ExtendedPayload)
  95                 default:
  96                         err = fmt.Errorf("unknown message type: %v", msg.Type)
  97                 }
  98         }
  99         data = make([]byte, 4+buf.Len())
 100         binary.BigEndian.PutUint32(data, uint32(buf.Len()))
 101         if buf.Len() != copy(data[4:], buf.Bytes()) {
 102                 panic("bad copy")
 103         }
 104         return
 105 }
 106
 107 type Decoder struct {
 108         R         *bufio.Reader
 109         MaxLength Integer // TODO: Should this include the length header or not?
 110 }
 111
 112 func (d *Decoder) Decode(msg *Message) (err error) {
 113         var length Integer
 114         err = binary.Read(d.R, binary.BigEndian, &length)
 115         if err != nil {
 116                 return
 117         }
 118         if length > d.MaxLength {
 119                 return errors.New("message too long")
 120         }
 121         if length == 0 {
 122                 msg.Keepalive = true
 123                 return
 124         }
 125         msg.Keepalive = false
 126         b := make([]byte, length)
 127         _, err = io.ReadFull(d.R, b)
 128         if err == io.EOF {
 129                 err = io.ErrUnexpectedEOF
 130                 return
 131         }
 132         if err != nil {
 133                 return
 134         }
 135         r := bytes.NewReader(b)
 136         defer func() {
 137                 written, _ := io.Copy(ioutil.Discard, r)
 138                 if written != 0 && err == nil {
 139                         err = fmt.Errorf("short read on message type %d, left %d bytes", msg.Type, written)
 140                 } else if err == io.EOF {
 141                         err = io.ErrUnexpectedEOF
 142                 }
 143         }()
 144         msg.Keepalive = false
 145         c, err := r.ReadByte()
 146         if err != nil {
 147                 return
 148         }
 149         msg.Type = MessageType(c)
 150         switch msg.Type {
 151         case Choke, Unchoke, Interested, NotInterested:
 152                 return
 153         case Have:
 154                 err = msg.Index.Read(r)
 155         case Request, Cancel:
 156                 for _, data := range []*Integer{&msg.Index, &msg.Begin, &msg.Length} {
 157                         err = data.Read(r)
 158                         if err != nil {
 159                                 break
 160                         }
 161                 }
 162         case Bitfield:
 163                 b := make([]byte, length-1)
 164                 _, err = io.ReadFull(r, b)
 165                 msg.Bitfield = unmarshalBitfield(b)
 166         case Piece:
 167                 for _, pi := range []*Integer{&msg.Index, &msg.Begin} {
 168                         err = pi.Read(r)
 169                         if err != nil {
 170                                 break
 171                         }
 172                 }
 173                 if err != nil {
 174                         break
 175                 }
 176                 msg.Piece, err = ioutil.ReadAll(r)
 177         case Extended:
 178                 msg.ExtendedID, err = r.ReadByte()
 179                 if err != nil {
 180                         break
 181                 }
 182                 msg.ExtendedPayload, err = ioutil.ReadAll(r)
 183         default:
 184                 err = fmt.Errorf("unknown message type %#v", c)
 185         }
 186         return
 187 }
 188
 189 type Bytes []byte
 190
 191 func (b Bytes) MarshalBinary() ([]byte, error) {
 192         return b, nil
 193 }
 194
 195 func unmarshalBitfield(b []byte) (bf []bool) {
 196         for _, c := range b {
 197                 for i := 7; i >= 0; i-- {
 198                         bf = append(bf, (c>>uint(i))&1 == 1)
 199                 }
 200         }
 201         return
 202 }
 203
 204 func marshalBitfield(bf []bool) (b []byte) {
 205         b = make([]byte, (len(bf)+7)/8)
 206         for i, have := range bf {
 207                 if !have {
 208                         continue
 209                 }
 210                 c := b[i/8]
 211                 c |= 1 << uint(7-i%8)
 212                 b[i/8] = c
 213         }
 214         return
 215 }