github.com/devwanda/aphelion-staking@v0.33.9/p2p/test_util.go (about) 1 package p2p 2 3 import ( 4 "fmt" 5 "net" 6 "time" 7 8 "github.com/pkg/errors" 9 10 "github.com/devwanda/aphelion-staking/crypto" 11 "github.com/devwanda/aphelion-staking/crypto/ed25519" 12 "github.com/devwanda/aphelion-staking/libs/log" 13 tmnet "github.com/devwanda/aphelion-staking/libs/net" 14 tmrand "github.com/devwanda/aphelion-staking/libs/rand" 15 16 "github.com/devwanda/aphelion-staking/config" 17 "github.com/devwanda/aphelion-staking/p2p/conn" 18 ) 19 20 const testCh = 0x01 21 22 //------------------------------------------------ 23 24 type mockNodeInfo struct { 25 addr *NetAddress 26 } 27 28 func (ni mockNodeInfo) ID() ID { return ni.addr.ID } 29 func (ni mockNodeInfo) NetAddress() (*NetAddress, error) { return ni.addr, nil } 30 func (ni mockNodeInfo) Validate() error { return nil } 31 func (ni mockNodeInfo) CompatibleWith(other NodeInfo) error { return nil } 32 33 func AddPeerToSwitchPeerSet(sw *Switch, peer Peer) { 34 sw.peers.Add(peer) 35 } 36 37 func CreateRandomPeer(outbound bool) Peer { 38 addr, netAddr := CreateRoutableAddr() 39 p := &peer{ 40 peerConn: peerConn{ 41 outbound: outbound, 42 socketAddr: netAddr, 43 }, 44 nodeInfo: mockNodeInfo{netAddr}, 45 mconn: &conn.MConnection{}, 46 metrics: NopMetrics(), 47 } 48 p.SetLogger(log.TestingLogger().With("peer", addr)) 49 return p 50 } 51 52 func CreateRoutableAddr() (addr string, netAddr *NetAddress) { 53 for { 54 var err error 55 addr = fmt.Sprintf("%X@%v.%v.%v.%v:26656", 56 tmrand.Bytes(20), 57 tmrand.Int()%256, 58 tmrand.Int()%256, 59 tmrand.Int()%256, 60 tmrand.Int()%256) 61 netAddr, err = NewNetAddressString(addr) 62 if err != nil { 63 panic(err) 64 } 65 if netAddr.Routable() { 66 break 67 } 68 } 69 return 70 } 71 72 //------------------------------------------------------------------ 73 // Connects switches via arbitrary net.Conn. Used for testing. 74 75 const TestHost = "localhost" 76 77 // MakeConnectedSwitches returns n switches, connected according to the connect func. 78 // If connect==Connect2Switches, the switches will be fully connected. 79 // initSwitch defines how the i'th switch should be initialized (ie. with what reactors). 80 // NOTE: panics if any switch fails to start. 81 func MakeConnectedSwitches(cfg *config.P2PConfig, 82 n int, 83 initSwitch func(int, *Switch) *Switch, 84 connect func([]*Switch, int, int), 85 ) []*Switch { 86 switches := make([]*Switch, n) 87 for i := 0; i < n; i++ { 88 switches[i] = MakeSwitch(cfg, i, TestHost, "123.123.123", initSwitch) 89 } 90 91 if err := StartSwitches(switches); err != nil { 92 panic(err) 93 } 94 95 for i := 0; i < n; i++ { 96 for j := i + 1; j < n; j++ { 97 connect(switches, i, j) 98 } 99 } 100 101 return switches 102 } 103 104 // Connect2Switches will connect switches i and j via net.Pipe(). 105 // Blocks until a connection is established. 106 // NOTE: caller ensures i and j are within bounds. 107 func Connect2Switches(switches []*Switch, i, j int) { 108 switchI := switches[i] 109 switchJ := switches[j] 110 111 c1, c2 := conn.NetPipe() 112 113 doneCh := make(chan struct{}) 114 go func() { 115 err := switchI.addPeerWithConnection(c1) 116 if err != nil { 117 panic(err) 118 } 119 doneCh <- struct{}{} 120 }() 121 go func() { 122 err := switchJ.addPeerWithConnection(c2) 123 if err != nil { 124 panic(err) 125 } 126 doneCh <- struct{}{} 127 }() 128 <-doneCh 129 <-doneCh 130 } 131 132 func (sw *Switch) addPeerWithConnection(conn net.Conn) error { 133 pc, err := testInboundPeerConn(conn, sw.config, sw.nodeKey.PrivKey) 134 if err != nil { 135 if err := conn.Close(); err != nil { 136 sw.Logger.Error("Error closing connection", "err", err) 137 } 138 return err 139 } 140 141 ni, err := handshake(conn, time.Second, sw.nodeInfo) 142 if err != nil { 143 if err := conn.Close(); err != nil { 144 sw.Logger.Error("Error closing connection", "err", err) 145 } 146 return err 147 } 148 149 p := newPeer( 150 pc, 151 MConnConfig(sw.config), 152 ni, 153 sw.reactorsByCh, 154 sw.chDescs, 155 sw.StopPeerForError, 156 ) 157 158 if err = sw.addPeer(p); err != nil { 159 pc.CloseConn() 160 return err 161 } 162 163 return nil 164 } 165 166 // StartSwitches calls sw.Start() for each given switch. 167 // It returns the first encountered error. 168 func StartSwitches(switches []*Switch) error { 169 for _, s := range switches { 170 err := s.Start() // start switch and reactors 171 if err != nil { 172 return err 173 } 174 } 175 return nil 176 } 177 178 func MakeSwitch( 179 cfg *config.P2PConfig, 180 i int, 181 network, version string, 182 initSwitch func(int, *Switch) *Switch, 183 opts ...SwitchOption, 184 ) *Switch { 185 186 nodeKey := NodeKey{ 187 PrivKey: ed25519.GenPrivKey(), 188 } 189 nodeInfo := testNodeInfo(nodeKey.ID(), fmt.Sprintf("node%d", i)) 190 addr, err := NewNetAddressString( 191 IDAddressString(nodeKey.ID(), nodeInfo.(DefaultNodeInfo).ListenAddr), 192 ) 193 if err != nil { 194 panic(err) 195 } 196 197 t := NewMultiplexTransport(nodeInfo, nodeKey, MConnConfig(cfg)) 198 199 if err := t.Listen(*addr); err != nil { 200 panic(err) 201 } 202 203 // TODO: let the config be passed in? 204 sw := initSwitch(i, NewSwitch(cfg, t, opts...)) 205 sw.SetLogger(log.TestingLogger().With("switch", i)) 206 sw.SetNodeKey(&nodeKey) 207 208 ni := nodeInfo.(DefaultNodeInfo) 209 for ch := range sw.reactorsByCh { 210 ni.Channels = append(ni.Channels, ch) 211 } 212 nodeInfo = ni 213 214 // TODO: We need to setup reactors ahead of time so the NodeInfo is properly 215 // populated and we don't have to do those awkward overrides and setters. 216 t.nodeInfo = nodeInfo 217 sw.SetNodeInfo(nodeInfo) 218 219 return sw 220 } 221 222 func testInboundPeerConn( 223 conn net.Conn, 224 config *config.P2PConfig, 225 ourNodePrivKey crypto.PrivKey, 226 ) (peerConn, error) { 227 return testPeerConn(conn, config, false, false, ourNodePrivKey, nil) 228 } 229 230 func testPeerConn( 231 rawConn net.Conn, 232 cfg *config.P2PConfig, 233 outbound, persistent bool, 234 ourNodePrivKey crypto.PrivKey, 235 socketAddr *NetAddress, 236 ) (pc peerConn, err error) { 237 conn := rawConn 238 239 // Fuzz connection 240 if cfg.TestFuzz { 241 // so we have time to do peer handshakes and get set up 242 conn = FuzzConnAfterFromConfig(conn, 10*time.Second, cfg.TestFuzzConfig) 243 } 244 245 // Encrypt connection 246 conn, err = upgradeSecretConn(conn, cfg.HandshakeTimeout, ourNodePrivKey) 247 if err != nil { 248 return pc, errors.Wrap(err, "Error creating peer") 249 } 250 251 // Only the information we already have 252 return newPeerConn(outbound, persistent, conn, socketAddr), nil 253 } 254 255 //---------------------------------------------------------------- 256 // rand node info 257 258 func testNodeInfo(id ID, name string) NodeInfo { 259 return testNodeInfoWithNetwork(id, name, "testing") 260 } 261 262 func testNodeInfoWithNetwork(id ID, name, network string) NodeInfo { 263 return DefaultNodeInfo{ 264 ProtocolVersion: defaultProtocolVersion, 265 DefaultNodeID: id, 266 ListenAddr: fmt.Sprintf("127.0.0.1:%d", getFreePort()), 267 Network: network, 268 Version: "1.2.3-rc0-deadbeef", 269 Channels: []byte{testCh}, 270 Moniker: name, 271 Other: DefaultNodeInfoOther{ 272 TxIndex: "on", 273 RPCAddress: fmt.Sprintf("127.0.0.1:%d", getFreePort()), 274 }, 275 } 276 } 277 278 func getFreePort() int { 279 port, err := tmnet.GetFreePort() 280 if err != nil { 281 panic(err) 282 } 283 return port 284 }