peer_protocol/pex.go

   1 package peer_protocol
   2
   3 import (
   4         "net"
   5
   6         "github.com/anacrolix/dht/v2/krpc"
   7         "github.com/anacrolix/torrent/bencode"
   8 )
   9
  10 type PexMsg struct {
  11         Added       krpc.CompactIPv4NodeAddrs `bencode:"added"`
  12         AddedFlags  []PexPeerFlags            `bencode:"added.f"`
  13         Added6      krpc.CompactIPv6NodeAddrs `bencode:"added6"`
  14         Added6Flags []PexPeerFlags            `bencode:"added6.f"`
  15         Dropped     krpc.CompactIPv4NodeAddrs `bencode:"dropped"`
  16         Dropped6    krpc.CompactIPv6NodeAddrs `bencode:"dropped6"`
  17 }
  18
  19 func addrEqual(a, b *krpc.NodeAddr) bool {
  20         return a.IP.Equal(b.IP) && a.Port == b.Port
  21 }
  22
  23 func addrIndex(v []krpc.NodeAddr, a *krpc.NodeAddr) int {
  24         for i := range v {
  25                 if addrEqual(&v[i], a) {
  26                         return i
  27                 }
  28         }
  29         return -1
  30 }
  31
  32 func (m *PexMsg) Add(addr krpc.NodeAddr, f PexPeerFlags) {
  33         if addr.IP.To4() != nil {
  34                 if addrIndex(m.Added.NodeAddrs(), &addr) >= 0 {
  35                         // already added
  36                         return
  37                 }
  38                 if i := addrIndex(m.Dropped.NodeAddrs(), &addr); i >= 0 {
  39                         // on the dropped list - cancel out
  40                         m.Dropped = append(m.Dropped[:i], m.Dropped[i+1:]...)
  41                         return
  42                 }
  43                 m.Added = append(m.Added, addr)
  44                 m.AddedFlags = append(m.AddedFlags, f)
  45         } else if len(addr.IP) == net.IPv6len {
  46                 if addrIndex(m.Added6.NodeAddrs(), &addr) >= 0 {
  47                         // already added
  48                         return
  49                 }
  50                 if i := addrIndex(m.Dropped6.NodeAddrs(), &addr); i >= 0 {
  51                         // on the dropped list - cancel out
  52                         m.Dropped6 = append(m.Dropped6[:i], m.Dropped6[i+1:]...)
  53                         return
  54                 }
  55                 m.Added6 = append(m.Added6, addr)
  56                 m.Added6Flags = append(m.Added6Flags, f)
  57         }
  58 }
  59
  60 func (m *PexMsg) Drop(addr krpc.NodeAddr) {
  61         if addr.IP.To4() != nil {
  62                 if addrIndex(m.Dropped.NodeAddrs(), &addr) >= 0 {
  63                         // already dropped
  64                         return
  65                 }
  66                 if i := addrIndex(m.Added.NodeAddrs(), &addr); i >= 0 {
  67                         // on the added list - cancel out
  68                         m.Added = append(m.Added[:i], m.Added[i+1:]...)
  69                         m.AddedFlags = append(m.AddedFlags[:i], m.AddedFlags[i+1:]...)
  70                         return
  71                 }
  72                 m.Dropped = append(m.Dropped, addr)
  73         } else if len(addr.IP) == net.IPv6len {
  74                 if addrIndex(m.Dropped6.NodeAddrs(), &addr) >= 0 {
  75                         // already dropped
  76                         return
  77                 }
  78                 if i := addrIndex(m.Added6.NodeAddrs(), &addr); i >= 0 {
  79                         // on the added list - cancel out
  80                         m.Added6 = append(m.Added6[:i], m.Added6[i+1:]...)
  81                         m.Added6Flags = append(m.Added6Flags[:i], m.Added6Flags[i+1:]...)
  82                         return
  83                 }
  84                 m.Dropped6 = append(m.Dropped6, addr)
  85         }
  86 }
  87
  88 func (m *PexMsg) Len() int {
  89         return len(m.Added)+len(m.Added6)+len(m.Dropped)+len(m.Dropped6)
  90 }
  91
  92 // DeltaLen returns max of {added+added6, dropped+dropped6}
  93 func (m *PexMsg) DeltaLen() int {
  94         lenAdded := len(m.Added)+len(m.Added6)
  95         lenDropped := len(m.Dropped)+len(m.Dropped6)
  96         if lenAdded > lenDropped {
  97                 return lenAdded
  98         }
  99         return lenDropped
 100 }
 101
 102 func (m *PexMsg) Message(pexExtendedId ExtensionNumber) Message {
 103         payload := bencode.MustMarshal(m)
 104         return Message{
 105                 Type:            Extended,
 106                 ExtendedID:      pexExtendedId,
 107                 ExtendedPayload: payload,
 108         }
 109 }
 110
 111 func LoadPexMsg(b []byte) (*PexMsg, error) {
 112         m := new(PexMsg)
 113         if err := bencode.Unmarshal(b, m); err != nil {
 114                 return nil, err
 115         }
 116         return m, nil
 117 }
 118
 119
 120 type PexPeerFlags byte
 121
 122 func (me PexPeerFlags) Get(f PexPeerFlags) bool {
 123         return me&f == f
 124 }
 125
 126 const (
 127         PexPrefersEncryption PexPeerFlags = 1 << iota
 128         PexSeedUploadOnly
 129         PexSupportsUtp
 130         PexHolepunchSupport
 131         PexOutgoingConn
 132 )