peer_protocol/msg.go

   1 package peer_protocol
   2
   3 import (
   4         "bytes"
   5         "encoding/binary"
   6         "fmt"
   7 )
   8
   9 // This is a lazy union representing all the possible fields for messages. Go doesn't have ADTs, and
  10 // I didn't choose to use type-assertions.
  11 type Message struct {
  12         Keepalive            bool
  13         Type                 MessageType
  14         Index, Begin, Length Integer
  15         Piece                []byte
  16         Bitfield             []bool
  17         ExtendedID           ExtensionNumber
  18         ExtendedPayload      []byte
  19         Port                 uint16
  20 }
  21
  22 func MakeCancelMessage(piece, offset, length Integer) Message {
  23         return Message{
  24                 Type:   Cancel,
  25                 Index:  piece,
  26                 Begin:  offset,
  27                 Length: length,
  28         }
  29 }
  30
  31 func (msg Message) RequestSpec() (ret RequestSpec) {
  32         return RequestSpec{
  33                 msg.Index,
  34                 msg.Begin,
  35                 func() Integer {
  36                         if msg.Type == Piece {
  37                                 return Integer(len(msg.Piece))
  38                         } else {
  39                                 return msg.Length
  40                         }
  41                 }(),
  42         }
  43 }
  44
  45 func (msg Message) MustMarshalBinary() []byte {
  46         b, err := msg.MarshalBinary()
  47         if err != nil {
  48                 panic(err)
  49         }
  50         return b
  51 }
  52
  53 func (msg Message) MarshalBinary() (data []byte, err error) {
  54         buf := &bytes.Buffer{}
  55         if !msg.Keepalive {
  56                 err = buf.WriteByte(byte(msg.Type))
  57                 if err != nil {
  58                         return
  59                 }
  60                 switch msg.Type {
  61                 case Choke, Unchoke, Interested, NotInterested, HaveAll, HaveNone:
  62                 case Have:
  63                         err = binary.Write(buf, binary.BigEndian, msg.Index)
  64                 case Request, Cancel, Reject:
  65                         for _, i := range []Integer{msg.Index, msg.Begin, msg.Length} {
  66                                 err = binary.Write(buf, binary.BigEndian, i)
  67                                 if err != nil {
  68                                         break
  69                                 }
  70                         }
  71                 case Bitfield:
  72                         _, err = buf.Write(marshalBitfield(msg.Bitfield))
  73                 case Piece:
  74                         for _, i := range []Integer{msg.Index, msg.Begin} {
  75                                 err = binary.Write(buf, binary.BigEndian, i)
  76                                 if err != nil {
  77                                         return
  78                                 }
  79                         }
  80                         n, err := buf.Write(msg.Piece)
  81                         if err != nil {
  82                                 break
  83                         }
  84                         if n != len(msg.Piece) {
  85                                 panic(n)
  86                         }
  87                 case Extended:
  88                         err = buf.WriteByte(byte(msg.ExtendedID))
  89                         if err != nil {
  90                                 return
  91                         }
  92                         _, err = buf.Write(msg.ExtendedPayload)
  93                 case Port:
  94                         err = binary.Write(buf, binary.BigEndian, msg.Port)
  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 func marshalBitfield(bf []bool) (b []byte) {
 108         b = make([]byte, (len(bf)+7)/8)
 109         for i, have := range bf {
 110                 if !have {
 111                         continue
 112                 }
 113                 c := b[i/8]
 114                 c |= 1 << uint(7-i%8)
 115                 b[i/8] = c
 116         }
 117         return
 118 }