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