forked from cenkalti/rain
-
Notifications
You must be signed in to change notification settings - Fork 0
/
addrlist.go
151 lines (135 loc) · 3.37 KB
/
addrlist.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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
package addrlist
import (
"net"
"sort"
"time"
"github.com/google/btree"
"github.com/kyoto44/rain/blocklist"
"github.com/kyoto44/rain/externalip"
"github.com/kyoto44/rain/peerpriority"
"github.com/kyoto44/rain/peersource"
)
// AddrList contains peer addresses that are ready to be connected.
type AddrList struct {
peerByTime []*peerAddr
peerByPriority *btree.BTree
maxItems int
listenPort int
clientIP *net.IP
blocklist *blocklist.Blocklist
countBySource map[peersource.Source]int
}
// New returns a new AddrList.
func New(maxItems int, blocklist *blocklist.Blocklist, listenPort int, clientIP *net.IP) *AddrList {
return &AddrList{
peerByPriority: btree.New(2),
maxItems: maxItems,
listenPort: listenPort,
clientIP: clientIP,
blocklist: blocklist,
countBySource: make(map[peersource.Source]int),
}
}
// Reset empties the address list.
func (d *AddrList) Reset() {
d.peerByTime = nil
d.peerByPriority.Clear(false)
d.countBySource = make(map[peersource.Source]int)
}
// Len returns the number of addresses in the list.
func (d *AddrList) Len() int {
return d.peerByPriority.Len()
}
// LenSource returns the number of addresses for a single source.
func (d *AddrList) LenSource(s peersource.Source) int {
return d.countBySource[s]
}
// Pop returns the next address. The returned address is removed from the list.
func (d *AddrList) Pop() (*net.TCPAddr, peersource.Source) {
item := d.peerByPriority.DeleteMax()
if item == nil {
return nil, 0
}
p := item.(*peerAddr)
d.peerByTime[p.index] = nil
d.countBySource[p.source]--
return p.addr, p.source
}
// Push adds a new address to the list. Does nothing if the address is already in the list.
func (d *AddrList) Push(addrs []*net.TCPAddr, source peersource.Source) {
now := time.Now()
var added int
for _, ad := range addrs {
// 0 port is invalid
if ad.Port == 0 {
continue
}
// Discard own client
if ad.IP.IsLoopback() && ad.Port == d.listenPort {
continue
} else if d.clientIP.Equal(ad.IP) {
continue
}
if externalip.IsExternal(ad.IP) {
continue
}
if d.blocklist != nil && d.blocklist.Blocked(ad.IP) {
continue
}
p := &peerAddr{
addr: ad,
timestamp: now,
source: source,
priority: peerpriority.Calculate(ad, d.clientAddr()),
}
item := d.peerByPriority.ReplaceOrInsert(p)
if item != nil {
prev := item.(*peerAddr)
d.peerByTime[prev.index] = p
p.index = prev.index
d.countBySource[prev.source]--
} else {
d.peerByTime = append(d.peerByTime, p)
p.index = len(d.peerByTime) - 1
}
added++
}
d.filterNils()
sort.Sort(byTimestamp(d.peerByTime))
d.countBySource[source] += added
delta := d.peerByPriority.Len() - d.maxItems
if delta > 0 {
d.removeExcessItems(delta)
d.filterNils()
d.countBySource[source] -= delta
}
if len(d.peerByTime) != d.peerByPriority.Len() {
panic("addr list data structures not in sync")
}
}
func (d *AddrList) filterNils() {
b := d.peerByTime[:0]
for _, x := range d.peerByTime {
if x != nil {
b = append(b, x)
x.index = len(b) - 1
}
}
d.peerByTime = b
}
func (d *AddrList) removeExcessItems(delta int) {
for i := 0; i < delta; i++ {
d.peerByPriority.Delete(d.peerByTime[i])
d.peerByTime[i] = nil
}
}
func (d *AddrList) clientAddr() *net.TCPAddr {
ip := *d.clientIP
if ip == nil {
ip = net.IPv4(0, 0, 0, 0)
}
return &net.TCPAddr{
IP: ip,
Port: d.listenPort,
}
}