-
Notifications
You must be signed in to change notification settings - Fork 199
/
hostWithConnectionManagement.go
60 lines (49 loc) · 1.58 KB
/
hostWithConnectionManagement.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
package discovery
import (
"context"
"fmt"
"github.com/ElrondNetwork/elrond-go/core/check"
"github.com/ElrondNetwork/elrond-go/p2p"
"github.com/libp2p/go-libp2p-core/peer"
)
type hostWithConnectionManagement struct {
sharder Sharder
ConnectableHost
}
// NewHostWithConnectionManagement returns a host wrapper able to decide if connection initiated to a peer
// will actually be kept or not
func NewHostWithConnectionManagement(ch ConnectableHost, sharder Sharder) (*hostWithConnectionManagement, error) {
if check.IfNil(ch) {
return nil, p2p.ErrNilHost
}
if check.IfNil(sharder) {
return nil, p2p.ErrNilSharder
}
return &hostWithConnectionManagement{
ConnectableHost: ch,
sharder: sharder,
}, nil
}
// Connect tries to connect to the provided address info if the sharder allows it
func (hwcm *hostWithConnectionManagement) Connect(ctx context.Context, pi peer.AddrInfo) error {
err := hwcm.canConnectToPeer(pi.ID)
if err != nil {
return err
}
return hwcm.ConnectableHost.Connect(ctx, pi)
}
func (hwcm *hostWithConnectionManagement) canConnectToPeer(pid peer.ID) error {
allPeers := hwcm.ConnectableHost.Network().Peers()
if !hwcm.sharder.Has(pid, allPeers) {
allPeers = append(allPeers, pid)
}
evicted := hwcm.sharder.ComputeEvictionList(allPeers)
if hwcm.sharder.Has(pid, evicted) {
return fmt.Errorf("%w, pid: %s", p2p.ErrUnwantedPeer, pid.Pretty())
}
return nil
}
// IsInterfaceNil returns true if there is no value under the interface
func (hwcm *hostWithConnectionManagement) IsInterfaceNil() bool {
return hwcm == nil || check.IfNil(hwcm.ConnectableHost)
}