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() (ret RequestSpec) {
  30         return RequestSpec{
  31                 msg.Index,
  32                 msg.Begin,
  33                 func() Integer {
  34                         if msg.Type == Piece {
  35                                 return Integer(len(msg.Piece))
  36                         } else {
  37                                 return msg.Length
  38                         }
  39                 }(),
  40         }
  41 }
  42
  43 func (msg Message) MustMarshalBinary() []byte {
  44         b, err := msg.MarshalBinary()
  45         if err != nil {
  46                 panic(err)
  47         }
  48         return b
  49 }
  50
  51 func (msg Message) MarshalBinary() (data []byte, err error) {
  52         buf := &bytes.Buffer{}
  53         if !msg.Keepalive {
  54                 err = buf.WriteByte(byte(msg.Type))
  55                 if err != nil {
  56                         return
  57                 }
  58                 switch msg.Type {
  59                 case Choke, Unchoke, Interested, NotInterested, HaveAll, HaveNone:
  60                 case Have:
  61                         err = binary.Write(buf, binary.BigEndian, msg.Index)
  62                 case Request, Cancel, Reject:
  63                         for _, i := range []Integer{msg.Index, msg.Begin, msg.Length} {
  64                                 err = binary.Write(buf, binary.BigEndian, i)
  65                                 if err != nil {
  66                                         break
  67                                 }
  68                         }
  69                 case Bitfield:
  70                         _, err = buf.Write(marshalBitfield(msg.Bitfield))
  71                 case Piece:
  72                         for _, i := range []Integer{msg.Index, msg.Begin} {
  73                                 err = binary.Write(buf, binary.BigEndian, i)
  74                                 if err != nil {
  75                                         return
  76                                 }
  77                         }
  78                         n, err := buf.Write(msg.Piece)
  79                         if err != nil {
  80                                 break
  81                         }
  82                         if n != len(msg.Piece) {
  83                                 panic(n)
  84                         }
  85                 case Extended:
  86                         err = buf.WriteByte(byte(msg.ExtendedID))
  87                         if err != nil {
  88                                 return
  89                         }
  90                         _, err = buf.Write(msg.ExtendedPayload)
  91                 case Port:
  92                         err = binary.Write(buf, binary.BigEndian, msg.Port)
  93                 default:
  94                         err = fmt.Errorf("unknown message type: %v", msg.Type)
  95                 }
  96         }
  97         data = make([]byte, 4+buf.Len())
  98         binary.BigEndian.PutUint32(data, uint32(buf.Len()))
  99         if buf.Len() != copy(data[4:], buf.Bytes()) {
 100                 panic("bad copy")
 101         }
 102         return
 103 }
 104
 105 func marshalBitfield(bf []bool) (b []byte) {
 106         b = make([]byte, (len(bf)+7)/8)
 107         for i, have := range bf {
 108                 if !have {
 109                         continue
 110                 }
 111                 c := b[i/8]
 112                 c |= 1 << uint(7-i%8)
 113                 b[i/8] = c
 114         }
 115         return
 116 }