peer_protocol/msg.go

   1 package peer_protocol
   2
   3 import (
   4         "bytes"
   5         "encoding/binary"
   6         "fmt"
   7 )
   8
   9 type Message struct {
  10         Keepalive            bool
  11         Type                 MessageType
  12         Index, Begin, Length Integer
  13         Piece                []byte
  14         Bitfield             []bool
  15         ExtendedID           ExtensionNumber
  16         ExtendedPayload      []byte
  17         Port                 uint16
  18 }
  19
  20 func MakeCancelMessage(piece, offset, length Integer) Message {
  21         return Message{
  22                 Type:   Cancel,
  23                 Index:  piece,
  24                 Begin:  offset,
  25                 Length: length,
  26         }
  27 }
  28
  29 func (msg Message) RequestSpec() RequestSpec {
  30         return RequestSpec{msg.Index, msg.Begin, msg.Length}
  31 }
  32
  33 func (msg Message) MustMarshalBinary() []byte {
  34         b, err := msg.MarshalBinary()
  35         if err != nil {
  36                 panic(err)
  37         }
  38         return b
  39 }
  40
  41 func (msg Message) MarshalBinary() (data []byte, err error) {
  42         buf := &bytes.Buffer{}
  43         if !msg.Keepalive {
  44                 err = buf.WriteByte(byte(msg.Type))
  45                 if err != nil {
  46                         return
  47                 }
  48                 switch msg.Type {
  49                 case Choke, Unchoke, Interested, NotInterested, HaveAll, HaveNone:
  50                 case Have:
  51                         err = binary.Write(buf, binary.BigEndian, msg.Index)
  52                 case Request, Cancel, Reject:
  53                         for _, i := range []Integer{msg.Index, msg.Begin, msg.Length} {
  54                                 err = binary.Write(buf, binary.BigEndian, i)
  55                                 if err != nil {
  56                                         break
  57                                 }
  58                         }
  59                 case Bitfield:
  60                         _, err = buf.Write(marshalBitfield(msg.Bitfield))
  61                 case Piece:
  62                         for _, i := range []Integer{msg.Index, msg.Begin} {
  63                                 err = binary.Write(buf, binary.BigEndian, i)
  64                                 if err != nil {
  65                                         return
  66                                 }
  67                         }
  68                         n, err := buf.Write(msg.Piece)
  69                         if err != nil {
  70                                 break
  71                         }
  72                         if n != len(msg.Piece) {
  73                                 panic(n)
  74                         }
  75                 case Extended:
  76                         err = buf.WriteByte(byte(msg.ExtendedID))
  77                         if err != nil {
  78                                 return
  79                         }
  80                         _, err = buf.Write(msg.ExtendedPayload)
  81                 case Port:
  82                         err = binary.Write(buf, binary.BigEndian, msg.Port)
  83                 default:
  84                         err = fmt.Errorf("unknown message type: %v", msg.Type)
  85                 }
  86         }
  87         data = make([]byte, 4+buf.Len())
  88         binary.BigEndian.PutUint32(data, uint32(buf.Len()))
  89         if buf.Len() != copy(data[4:], buf.Bytes()) {
  90                 panic("bad copy")
  91         }
  92         return
  93 }
  94
  95 func marshalBitfield(bf []bool) (b []byte) {
  96         b = make([]byte, (len(bf)+7)/8)
  97         for i, have := range bf {
  98                 if !have {
  99                         continue
 100                 }
 101                 c := b[i/8]
 102                 c |= 1 << uint(7-i%8)
 103                 b[i/8] = c
 104         }
 105         return
 106 }