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 var PanicOnAddPeerErr = false
  14
  15 // Switches connected via arbitrary net.Conn; useful for testing
  16 // Returns n switches, connected according to the connect func.
  17 // If connect==Connect2Switches, the switches will be fully connected.
  18 // initSwitch defines how the ith switch should be initialized (ie. with what reactors).
  19 // NOTE: panics if any switch fails to start.
  20 func MakeConnectedSwitches(cfg *cfg.P2PConfig, n int, initSwitch func(int, *Switch) *Switch, connect func([]*Switch, int, int)) []*Switch {
  21         switches := make([]*Switch, n)
  22         for i := 0; i < n; i++ {
  23                 switches[i] = makeSwitch(cfg, i, "testing", "123.123.123", initSwitch)
  24         }
  25
  26         if err := startSwitches(switches); err != nil {
  27                 panic(err)
  28         }
  29
  30         for i := 0; i < n; i++ {
  31                 for j := i; j < n; j++ {
  32                         connect(switches, i, j)
  33                 }
  34         }
  35
  36         return switches
  37 }
  38
  39 // Will connect switches i and j via net.Pipe()
  40 // Blocks until a conection is established.
  41 // NOTE: caller ensures i and j are within bounds
  42 func Connect2Switches(switches []*Switch, i, j int) {
  43         switchI := switches[i]
  44         switchJ := switches[j]
  45         c1, c2 := net.Pipe()
  46         doneCh := make(chan struct{})
  47         go func() {
  48                 err := switchI.addPeerWithConnection(c1)
  49                 if PanicOnAddPeerErr && err != nil {
  50                         panic(err)
  51                 }
  52                 doneCh <- struct{}{}
  53         }()
  54         go func() {
  55                 err := switchJ.addPeerWithConnection(c2)
  56                 if PanicOnAddPeerErr && err != nil {
  57                         panic(err)
  58                 }
  59                 doneCh <- struct{}{}
  60         }()
  61         <-doneCh
  62         <-doneCh
  63 }
  64
  65 func startSwitches(switches []*Switch) error {
  66         for _, s := range switches {
  67                 _, err := s.Start() // start switch and reactors
  68                 if err != nil {
  69                         return err
  70                 }
  71         }
  72         return nil
  73 }
  74
  75 func makeSwitch(cfg *cfg.P2PConfig, i int, network, version string, initSwitch func(int, *Switch) *Switch) *Switch {
  76         privKey := crypto.GenPrivKeyEd25519()
  77         // new switch, add reactors
  78         // TODO: let the config be passed in?
  79         s := initSwitch(i, NewSwitch(cfg, nil))
  80         s.SetNodeInfo(&NodeInfo{
  81                 PubKey:     privKey.PubKey().Unwrap().(crypto.PubKeyEd25519),
  82                 Moniker:    cmn.Fmt("switch%d", i),
  83                 Network:    network,
  84                 Version:    version,
  85                 RemoteAddr: cmn.Fmt("%v:%v", network, rand.Intn(64512)+1023),
  86                 ListenAddr: cmn.Fmt("%v:%v", network, rand.Intn(64512)+1023),
  87         })
  88         s.SetNodePrivKey(privKey)
  89         return s
  90 }