util/types.go

   1 package util
   2
   3 import (
   4         "encoding"
   5         "encoding/binary"
   6         "errors"
   7         "fmt"
   8         "net"
   9
  10         "github.com/bradfitz/iter"
  11
  12         "github.com/anacrolix/torrent/bencode"
  13 )
  14
  15 // Concatenated 6-byte peer addresses.
  16 type CompactIPv4Peers []CompactPeer
  17
  18 var (
  19         // This allows bencode.Unmarshal to do better than a string or []byte.
  20         _ bencode.Unmarshaler      = &CompactIPv4Peers{}
  21         _ encoding.BinaryMarshaler = CompactIPv4Peers{}
  22 )
  23
  24 // This allows bencode.Unmarshal to do better than a string or []byte.
  25 func (me *CompactIPv4Peers) UnmarshalBencode(b []byte) (err error) {
  26         var bb []byte
  27         err = bencode.Unmarshal(b, &bb)
  28         if err != nil {
  29                 return
  30         }
  31         *me, err = UnmarshalIPv4CompactPeers(bb)
  32         return
  33 }
  34
  35 func (me CompactIPv4Peers) MarshalBinary() (ret []byte, err error) {
  36         ret = make([]byte, len(me)*6)
  37         for i, cp := range me {
  38                 copy(ret[6*i:], cp.IP.To4())
  39                 binary.BigEndian.PutUint16(ret[6*i+4:], uint16(cp.Port))
  40         }
  41         return
  42 }
  43
  44 // Represents peer address in either IPv6 or IPv4 form.
  45 type CompactPeer struct {
  46         IP   net.IP
  47         Port int
  48 }
  49
  50 var (
  51         _ bencode.Marshaler   = &CompactPeer{}
  52         _ bencode.Unmarshaler = &CompactPeer{}
  53 )
  54
  55 func (me CompactPeer) MarshalBencode() (ret []byte, err error) {
  56         ip := me.IP
  57         if ip4 := ip.To4(); ip4 != nil {
  58                 ip = ip4
  59         }
  60         ret = make([]byte, len(ip)+2)
  61         copy(ret, ip)
  62         binary.BigEndian.PutUint16(ret[len(ip):], uint16(me.Port))
  63         return bencode.Marshal(ret)
  64 }
  65
  66 func (me *CompactPeer) UnmarshalBinary(b []byte) error {
  67         switch len(b) {
  68         case 18:
  69                 me.IP = make([]byte, 16)
  70         case 6:
  71                 me.IP = make([]byte, 4)
  72         default:
  73                 return fmt.Errorf("bad compact peer string: %q", b)
  74         }
  75         copy(me.IP, b)
  76         b = b[len(me.IP):]
  77         me.Port = int(binary.BigEndian.Uint16(b))
  78         return nil
  79 }
  80
  81 func (me *CompactPeer) UnmarshalBencode(b []byte) (err error) {
  82         var _b []byte
  83         err = bencode.Unmarshal(b, &_b)
  84         if err != nil {
  85                 return
  86         }
  87         return me.UnmarshalBinary(_b)
  88 }
  89
  90 func UnmarshalIPv4CompactPeers(b []byte) (ret []CompactPeer, err error) {
  91         if len(b)%6 != 0 {
  92                 err = errors.New("bad length")
  93                 return
  94         }
  95         num := len(b) / 6
  96         ret = make([]CompactPeer, num)
  97         for i := range iter.N(num) {
  98                 off := i * 6
  99                 err = ret[i].UnmarshalBinary(b[off : off+6])
 100                 if err != nil {
 101                         return
 102                 }
 103         }
 104         return
 105 }