gitlab.com/SiaPrime/SiaPrime@v1.4.1/modules/renter/hostdb/hosttree/hosttree.go (about) 1 package hosttree 2 3 import ( 4 "sort" 5 "sync" 6 7 "gitlab.com/NebulousLabs/errors" 8 "gitlab.com/NebulousLabs/fastrand" 9 10 "gitlab.com/SiaPrime/SiaPrime/build" 11 "gitlab.com/SiaPrime/SiaPrime/modules" 12 "gitlab.com/SiaPrime/SiaPrime/types" 13 ) 14 15 var ( 16 // ErrHostExists is returned if an Insert is called with a public key that 17 // already exists in the tree. 18 ErrHostExists = errors.New("host already exists in the tree") 19 20 // errNegativeWeight is returned from an Insert() call if an entry with a 21 // negative weight is added to the tree. Entries must always have a positive 22 // weight. 23 errNegativeWeight = errors.New("cannot insert using a negative weight") 24 25 // errNilEntry is returned if a fetch call results in a nil tree entry. nodes 26 // should always have a non-nil entry, unless they have been Delete()ed. 27 errNilEntry = errors.New("node has a nil entry") 28 29 // ErrNoSuchHost is returned if Remove is called with a public key that does 30 // not exist in the tree. 31 ErrNoSuchHost = errors.New("no host with specified public key") 32 33 // errWeightTooHeavy is returned from a SelectRandom() call if a weight that exceeds 34 // the total weight of the tree is requested. 35 errWeightTooHeavy = errors.New("requested a too-heavy weight") 36 ) 37 38 type ( 39 // WeightFunc is a function used to weight a given HostDBEntry in the tree. 40 WeightFunc func(modules.HostDBEntry) ScoreBreakdown 41 42 // HostTree is used to store and select host database entries. Each HostTree 43 // is initialized with a weighting func that is able to assign a weight to 44 // each entry. The entries can then be selected at random, weighted by the 45 // weight func. 46 HostTree struct { 47 root *node 48 49 // hosts is a map of public keys to nodes. 50 hosts map[string]*node 51 52 // resolver is the Resolver that is used by the hosttree to resolve 53 // hostnames to IP addresses. 54 resolver modules.Resolver 55 56 // weightFn calculates the weight of a hostEntry 57 weightFn WeightFunc 58 59 //filterByIP switches the IP Address subnet filtering on or off 60 filterByIP bool 61 62 mu sync.Mutex 63 } 64 65 // hostEntry is an entry in the host tree. 66 hostEntry struct { 67 modules.HostDBEntry 68 weight types.Currency 69 } 70 71 // node is a node in the tree. 72 node struct { 73 parent *node 74 left *node 75 right *node 76 77 count int // cumulative count of this node and all children 78 taken bool // `taken` indicates whether there is an active host at this node or not. 79 80 weight types.Currency 81 entry *hostEntry 82 } 83 ) 84 85 // createNode creates a new node using the provided `parent` and `entry`. 86 func createNode(parent *node, entry *hostEntry) *node { 87 return &node{ 88 parent: parent, 89 weight: entry.weight, 90 count: 1, 91 92 taken: true, 93 entry: entry, 94 } 95 } 96 97 // New creates a new HostTree given a weight function and a resolver 98 // for hostnames. 99 func New(wf WeightFunc, resolver modules.Resolver) *HostTree { 100 return &HostTree{ 101 hosts: make(map[string]*node), 102 root: &node{ 103 count: 1, 104 }, 105 resolver: resolver, 106 weightFn: wf, 107 filterByIP: false, 108 // NOTE: IPSubnetFiltering disabled by default 109 } 110 } 111 112 // recursiveInsert inserts an entry into the appropriate place in the tree. The 113 // running time of recursiveInsert is log(n) in the maximum number of elements 114 // that have ever been in the tree. 115 func (n *node) recursiveInsert(entry *hostEntry) (nodesAdded int, newnode *node) { 116 // If there is no parent and no children, and the node is not taken, assign 117 // this entry to this node. 118 if n.parent == nil && n.left == nil && n.right == nil && !n.taken { 119 n.entry = entry 120 n.taken = true 121 n.weight = entry.weight 122 newnode = n 123 return 124 } 125 126 n.weight = n.weight.Add(entry.weight) 127 128 // If the current node is empty, add the entry but don't increase the 129 // count. 130 if !n.taken { 131 n.taken = true 132 n.entry = entry 133 newnode = n 134 return 135 } 136 137 // Insert the element into the lest populated side. 138 if n.left == nil { 139 n.left = createNode(n, entry) 140 nodesAdded = 1 141 newnode = n.left 142 } else if n.right == nil { 143 n.right = createNode(n, entry) 144 nodesAdded = 1 145 newnode = n.right 146 } else if n.left.count <= n.right.count { 147 nodesAdded, newnode = n.left.recursiveInsert(entry) 148 } else { 149 nodesAdded, newnode = n.right.recursiveInsert(entry) 150 } 151 152 n.count += nodesAdded 153 return 154 } 155 156 // nodeAtWeight grabs an element in the tree that appears at the given weight. 157 // Though the tree has an arbitrary sorting, a sufficiently random weight will 158 // pull a random element. The tree is searched through in a post-ordered way. 159 func (n *node) nodeAtWeight(weight types.Currency) *node { 160 // Sanity check - weight must be less than the total weight of the tree. 161 if weight.Cmp(n.weight) > 0 { 162 build.Critical("Node weight corruption") 163 return nil 164 } 165 166 // Check if the left or right child should be returned. 167 if n.left != nil { 168 if weight.Cmp(n.left.weight) < 0 { 169 return n.left.nodeAtWeight(weight) 170 } 171 weight = weight.Sub(n.left.weight) // Search from the 0th index of the right side. 172 } 173 if n.right != nil && weight.Cmp(n.right.weight) < 0 { 174 return n.right.nodeAtWeight(weight) 175 } 176 177 // Should we panic here instead? 178 if !n.taken { 179 build.Critical("Node tree structure corruption") 180 return nil 181 } 182 183 // Return the root entry. 184 return n 185 } 186 187 // remove takes a node and removes it from the tree by climbing through the 188 // list of parents. remove does not delete nodes. 189 func (n *node) remove() { 190 n.weight = n.weight.Sub(n.entry.weight) 191 n.taken = false 192 current := n.parent 193 for current != nil { 194 current.weight = current.weight.Sub(n.entry.weight) 195 current = current.parent 196 } 197 } 198 199 // Host returns the address of the HostEntry. 200 func (he *hostEntry) Host() string { 201 return he.NetAddress.Host() 202 } 203 204 // All returns all of the hosts in the host tree, sorted by weight. 205 func (ht *HostTree) All() []modules.HostDBEntry { 206 ht.mu.Lock() 207 allHosts := ht.all() 208 ht.mu.Unlock() 209 return allHosts 210 } 211 212 // Insert inserts the entry provided to `entry` into the host tree. Insert will 213 // return an error if the input host already exists. 214 func (ht *HostTree) Insert(hdbe modules.HostDBEntry) error { 215 ht.mu.Lock() 216 defer ht.mu.Unlock() 217 return ht.insert(hdbe) 218 } 219 220 // Remove removes the host with the public key provided by `pk`. 221 func (ht *HostTree) Remove(pk types.SiaPublicKey) error { 222 ht.mu.Lock() 223 defer ht.mu.Unlock() 224 225 node, exists := ht.hosts[pk.String()] 226 if !exists { 227 return ErrNoSuchHost 228 } 229 node.remove() 230 delete(ht.hosts, pk.String()) 231 232 return nil 233 } 234 235 // Modify updates a host entry at the given public key, replacing the old entry 236 // with the entry provided by `newEntry`. 237 func (ht *HostTree) Modify(hdbe modules.HostDBEntry) error { 238 ht.mu.Lock() 239 defer ht.mu.Unlock() 240 241 node, exists := ht.hosts[hdbe.PublicKey.String()] 242 if !exists { 243 return ErrNoSuchHost 244 } 245 246 node.remove() 247 248 entry := &hostEntry{ 249 HostDBEntry: hdbe, 250 weight: ht.weightFn(hdbe).Score(), 251 } 252 253 _, node = ht.root.recursiveInsert(entry) 254 255 ht.hosts[entry.PublicKey.String()] = node 256 return nil 257 } 258 259 // SetFiltered updates a host entry filtered field. 260 func (ht *HostTree) SetFiltered(pubKey types.SiaPublicKey, filtered bool) error { 261 entry, ok := ht.Select(pubKey) 262 if !ok { 263 return ErrNoSuchHost 264 } 265 entry.Filtered = filtered 266 return ht.Modify(entry) 267 } 268 269 // SetWeightFunction resets the HostTree and assigns it a new weight 270 // function. This resets the tree and reinserts all the hosts. 271 func (ht *HostTree) SetWeightFunction(wf WeightFunc) error { 272 ht.mu.Lock() 273 defer ht.mu.Unlock() 274 275 // Get all the hosts. 276 allHosts := ht.all() 277 278 // Reset the tree 279 ht.hosts = make(map[string]*node) 280 ht.root = &node{ 281 count: 1, 282 } 283 284 // Assign the new weight function. 285 ht.weightFn = wf 286 287 // Reinsert all the hosts. To prevent the host tree from having a 288 // catastrophic failure in the event of an error early on, we tally up all 289 // of the insertion errors and return them all at the end. 290 var insertErrs error 291 for _, hdbe := range allHosts { 292 if err := ht.insert(hdbe); err != nil { 293 insertErrs = errors.Compose(err, insertErrs) 294 } 295 } 296 return insertErrs 297 } 298 299 // Select returns the host with the provided public key, should the host exist. 300 func (ht *HostTree) Select(spk types.SiaPublicKey) (modules.HostDBEntry, bool) { 301 ht.mu.Lock() 302 defer ht.mu.Unlock() 303 304 node, exists := ht.hosts[spk.String()] 305 if !exists { 306 return modules.HostDBEntry{}, false 307 } 308 return node.entry.HostDBEntry, true 309 } 310 311 // SelectRandom grabs a random n hosts from the tree. There will be no repeats, 312 // but the length of the slice returned may be less than n, and may even be 313 // zero. The hosts that are returned first have the higher priority. 314 // 315 // Hosts passed to 'blacklist' will not be considered; pass `nil` if no 316 // blacklist is desired. 'addressBlacklist' is similar to 'blacklist' but 317 // instead of not considering the hosts in the list, hosts that use the same IP 318 // subnet as those hosts will be ignored. In most cases those blacklists contain 319 // the same elements but sometimes it is useful to block a host without blocking 320 // its IP range. 321 // 322 // Hosts with a score of 1 will be ignored. 1 is the lowest score possible, at 323 // which point it's impossible to distinguish between hosts. Any sane scoring 324 // system should always have scores greater than 1 unless the host is 325 // intentionally being given a low score to indicate that the host should not be 326 // used. 327 func (ht *HostTree) SelectRandom(n int, blacklist, addressBlacklist []types.SiaPublicKey) []modules.HostDBEntry { 328 ht.mu.Lock() 329 defer ht.mu.Unlock() 330 331 var removedEntries []*hostEntry 332 333 // Create a filter. 334 filter := NewFilter(ht.resolver) 335 336 // Add the hosts from the addressBlacklist to the filter. 337 for _, pubkey := range addressBlacklist { 338 node, exists := ht.hosts[pubkey.String()] 339 if !exists { 340 continue 341 } 342 // Add the node to the addressFilter. 343 filter.Add(node.entry.NetAddress) 344 } 345 // Remove hosts we want to blacklist from the tree but remember them to make 346 // sure we can insert them later. 347 for _, pubkey := range blacklist { 348 node, exists := ht.hosts[pubkey.String()] 349 if !exists { 350 continue 351 } 352 // Remove the host from the tree. 353 node.remove() 354 delete(ht.hosts, pubkey.String()) 355 356 // Remember the host to insert it again later. 357 removedEntries = append(removedEntries, node.entry) 358 } 359 360 var hosts []modules.HostDBEntry 361 362 for len(hosts) < n && len(ht.hosts) > 0 { 363 randWeight := fastrand.BigIntn(ht.root.weight.Big()) 364 node := ht.root.nodeAtWeight(types.NewCurrency(randWeight)) 365 weightOne := types.NewCurrency64(1) 366 367 if node.entry.AcceptingContracts && 368 len(node.entry.ScanHistory) > 0 && 369 node.entry.ScanHistory[len(node.entry.ScanHistory)-1].Success && 370 !filter.Filtered(node.entry.NetAddress) && 371 node.entry.weight.Cmp(weightOne) > 0 { 372 // The host must be online and accepting contracts to be returned 373 // by the random function. It also has to pass the addressFilter 374 // check. 375 hosts = append(hosts, node.entry.HostDBEntry) 376 377 // If the host passed the filter, we add it to the filter if 378 // the CheckForIPViolation switch is set to true so the filter will disable 379 // selecting hosts with NetAddress from the same subnet 380 if ht.filterByIP { 381 filter.Add(node.entry.NetAddress) 382 } 383 } 384 385 removedEntries = append(removedEntries, node.entry) 386 node.remove() 387 delete(ht.hosts, node.entry.PublicKey.String()) 388 } 389 390 for _, entry := range removedEntries { 391 _, node := ht.root.recursiveInsert(entry) 392 ht.hosts[entry.PublicKey.String()] = node 393 } 394 395 return hosts 396 } 397 398 // all returns all of the hosts in the host tree, sorted by weight. 399 func (ht *HostTree) all() []modules.HostDBEntry { 400 var he []hostEntry 401 for _, node := range ht.hosts { 402 he = append(he, *node.entry) 403 } 404 sort.Sort(byWeight(he)) 405 406 var entries []modules.HostDBEntry 407 for _, entry := range he { 408 entries = append(entries, entry.HostDBEntry) 409 } 410 return entries 411 } 412 413 // insert inserts the entry provided to `entry` into the host tree. Insert will 414 // return an error if the input host already exists. 415 func (ht *HostTree) insert(hdbe modules.HostDBEntry) error { 416 entry := &hostEntry{ 417 HostDBEntry: hdbe, 418 weight: ht.weightFn(hdbe).Score(), 419 } 420 421 if _, exists := ht.hosts[entry.PublicKey.String()]; exists { 422 return ErrHostExists 423 } 424 425 _, node := ht.root.recursiveInsert(entry) 426 427 ht.hosts[entry.PublicKey.String()] = node 428 return nil 429 } 430 431 // SetFilterByIPEnabled enables or disables the host filtering by IP address. 432 // If enabled HostTree does not select more than one host from a single IP or 433 // the same subnet 434 func (ht *HostTree) SetFilterByIPEnabled(enabled bool) { 435 ht.mu.Lock() 436 ht.filterByIP = enabled 437 ht.mu.Unlock() 438 } 439 440 // FilterByIPEnabled tells if the IP filtering (Same subnet IP) is enabled or not 441 func (ht *HostTree) FilterByIPEnabled() bool { 442 ht.mu.Lock() 443 defer ht.mu.Unlock() 444 return ht.filterByIP 445 }