p2p/test_util.go

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