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