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