p2p/test_util.go

   1 package p2p
   2
   3 import (
   4         "math/rand"
   5         "net"
   6
   7         "github.com/tendermint/go-crypto"
   8         cmn "github.com/tendermint/tmlibs/common"
   9
  10         cfg "github.com/bytom/config"
  11 )
  12
  13 //PanicOnAddPeerErr add peer error
  14 var PanicOnAddPeerErr = false
  15
  16 func CreateRandomPeer(outbound bool) *Peer {
  17         _, netAddr := CreateRoutableAddr()
  18         p := &Peer{
  19                 peerConn: &peerConn{
  20                         outbound: outbound,
  21                 },
  22                 NodeInfo: &NodeInfo{
  23                         ListenAddr: netAddr.DialString(),
  24                 },
  25                 mconn: &MConnection{},
  26         }
  27         return p
  28 }
  29
  30 func CreateRoutableAddr() (addr string, netAddr *NetAddress) {
  31         for {
  32                 var err error
  33                 addr = cmn.Fmt("%X@%v.%v.%v.%v:46656", cmn.RandBytes(20), cmn.RandInt()%256, cmn.RandInt()%256, cmn.RandInt()%256, cmn.RandInt()%256)
  34                 netAddr, err = NewNetAddressString(addr)
  35                 if err != nil {
  36                         panic(err)
  37                 }
  38                 if netAddr.Routable() {
  39                         break
  40                 }
  41         }
  42         return
  43 }
  44
  45 // MakeConnectedSwitches switches connected via arbitrary net.Conn; useful for testing
  46 // Returns n switches, connected according to the connect func.
  47 // If connect==Connect2Switches, the switches will be fully connected.
  48 // initSwitch defines how the ith switch should be initialized (ie. with what reactors).
  49 // NOTE: panics if any switch fails to start.
  50 func MakeConnectedSwitches(cfg *cfg.P2PConfig, n int, initSwitch func(int, *Switch) *Switch, connect func([]*Switch, int, int)) []*Switch {
  51         switches := make([]*Switch, n)
  52         for i := 0; i < n; i++ {
  53                 switches[i] = MakeSwitch(cfg, i, "testing", "123.123.123", initSwitch)
  54         }
  55
  56         if err := startSwitches(switches); err != nil {
  57                 panic(err)
  58         }
  59
  60         for i := 0; i < n; i++ {
  61                 for j := i; j < n; j++ {
  62                         connect(switches, i, j)
  63                 }
  64         }
  65
  66         return switches
  67 }
  68
  69 // Connect2Switches will connect switches i and j via net.Pipe()
  70 // Blocks until a conection is established.
  71 // NOTE: caller ensures i and j are within bounds
  72 func Connect2Switches(switches []*Switch, i, j int) {
  73         switchI := switches[i]
  74         switchJ := switches[j]
  75         c1, c2 := net.Pipe()
  76         doneCh := make(chan struct{})
  77         go func() {
  78                 err := switchI.addPeerWithConnection(c1)
  79                 if PanicOnAddPeerErr && err != nil {
  80                         panic(err)
  81                 }
  82                 doneCh <- struct{}{}
  83         }()
  84         go func() {
  85                 err := switchJ.addPeerWithConnection(c2)
  86                 if PanicOnAddPeerErr && err != nil {
  87                         panic(err)
  88                 }
  89                 doneCh <- struct{}{}
  90         }()
  91         <-doneCh
  92         <-doneCh
  93 }
  94
  95 func startSwitches(switches []*Switch) error {
  96         for _, s := range switches {
  97                 _, err := s.Start() // start switch and reactors
  98                 if err != nil {
  99                         return err
 100                 }
 101         }
 102         return nil
 103 }
 104
 105 func MakeSwitch(cfg *cfg.P2PConfig, i int, network, version string, initSwitch func(int, *Switch) *Switch) *Switch {
 106         privKey := crypto.GenPrivKeyEd25519()
 107         // new switch, add reactors
 108         // TODO: let the config be passed in?
 109         s := initSwitch(i, NewSwitch(cfg, nil, nil))
 110         s.SetNodeInfo(&NodeInfo{
 111                 PubKey:     privKey.PubKey().Unwrap().(crypto.PubKeyEd25519),
 112                 Moniker:    cmn.Fmt("switch%d", i),
 113                 Network:    network,
 114                 Version:    version,
 115                 RemoteAddr: cmn.Fmt("%v:%v", network, rand.Intn(64512)+1023),
 116                 ListenAddr: cmn.Fmt("%v:%v", network, rand.Intn(64512)+1023),
 117         })
 118         s.SetNodePrivKey(privKey)
 119         return s
 120 }