p2p/peer_test.go

   1 package p2p
   2
   3 import (
   4         "fmt"
   5         "net"
   6         "testing"
   7         "time"
   8
   9         cfg "github.com/vapor/config"
  10         "github.com/vapor/consensus"
  11         conn "github.com/vapor/p2p/connection"
  12         "github.com/vapor/p2p/signlib"
  13         "github.com/vapor/version"
  14 )
  15
  16 const testCh = 0x01
  17
  18 func TestPeerBasic(t *testing.T) {
  19         privKey, err := signlib.NewPrivKey()
  20         if err != nil {
  21                 t.Fatal(err)
  22         }
  23
  24         // simulate remote peer
  25         rp := &remotePeer{PrivKey: privKey, Config: testCfg}
  26         rp.Start()
  27         defer rp.Stop()
  28
  29         p, err := createOutboundPeerAndPerformHandshake(rp.Addr(), cfg.DefaultP2PConfig())
  30         if err != nil {
  31                 t.Fatal(err)
  32         }
  33         _, err = p.Start()
  34         if err != nil {
  35                 t.Fatal(err)
  36         }
  37         defer p.Stop()
  38 }
  39
  40 func TestPeerSend(t *testing.T) {
  41         config := testCfg
  42         privKey, err := signlib.NewPrivKey()
  43         if err != nil {
  44                 t.Fatal(err)
  45         }
  46         // simulate remote peer
  47         rp := &remotePeer{PrivKey: privKey, Config: config}
  48         rp.Start()
  49         defer rp.Stop()
  50
  51         p, err := createOutboundPeerAndPerformHandshake(rp.Addr(), config.P2P)
  52         if err != nil {
  53                 t.Fatal(err)
  54         }
  55
  56         _, err = p.Start()
  57         if err != nil {
  58                 t.Fatal(err)
  59         }
  60
  61         defer p.Stop()
  62         if ok := p.CanSend(testCh); !ok {
  63                 t.Fatal("TestPeerSend send err")
  64         }
  65
  66         if ok := p.TrySend(testCh, []byte("test date")); !ok {
  67                 t.Fatal("TestPeerSend try send err")
  68         }
  69 }
  70
  71 func createOutboundPeerAndPerformHandshake(
  72         addr *NetAddress,
  73         config *cfg.P2PConfig,
  74 ) (*Peer, error) {
  75         chDescs := []*conn.ChannelDescriptor{
  76                 {ID: testCh, Priority: 1},
  77         }
  78         reactorsByCh := map[byte]Reactor{testCh: NewTestReactor(chDescs, true)}
  79         privKey, err := signlib.NewPrivKey()
  80         if err != nil {
  81                 return nil, err
  82         }
  83         peerConfig := DefaultPeerConfig(config)
  84         pc, err := newOutboundPeerConn(addr, privKey, peerConfig)
  85         if err != nil {
  86                 return nil, err
  87         }
  88         nodeInfo, err := pc.HandshakeTimeout(&NodeInfo{
  89                 Moniker: "host_peer",
  90                 Network: "testing",
  91                 Version: "123.123.123",
  92         }, 5*time.Second)
  93         if err != nil {
  94                 fmt.Println(err)
  95                 return nil, err
  96         }
  97         p := newPeer(pc, nodeInfo, reactorsByCh, chDescs, nil, false)
  98         return p, nil
  99 }
 100
 101 type remotePeer struct {
 102         PrivKey    signlib.PrivKey
 103         Config     *cfg.Config
 104         addr       *NetAddress
 105         quit       chan struct{}
 106         listenAddr string
 107 }
 108
 109 func (rp *remotePeer) Addr() *NetAddress {
 110         return rp.addr
 111 }
 112
 113 func (rp *remotePeer) Start() {
 114         if rp.listenAddr == "" {
 115                 rp.listenAddr = "127.0.0.1:0"
 116         }
 117
 118         l, e := net.Listen("tcp", rp.listenAddr) // any available address
 119         if e != nil {
 120                 fmt.Println("net.Listen tcp :0:", e)
 121         }
 122         rp.addr = NewNetAddress(l.Addr())
 123         rp.quit = make(chan struct{})
 124         go rp.accept(l)
 125 }
 126
 127 func (rp *remotePeer) Stop() {
 128         close(rp.quit)
 129 }
 130
 131 func (rp *remotePeer) accept(l net.Listener) {
 132         conns := []net.Conn{}
 133
 134         for {
 135                 conn, err := l.Accept()
 136                 if err != nil {
 137                         fmt.Println("Failed to accept conn:", err)
 138                 }
 139
 140                 pc, err := newInboundPeerConn(conn, rp.PrivKey, rp.Config.P2P)
 141                 if err != nil {
 142                         fmt.Println("Failed to create a peer:", err)
 143                 }
 144
 145                 _, err = pc.HandshakeTimeout(&NodeInfo{
 146                         PubKey:      rp.PrivKey.XPub().String(),
 147                         Moniker:     "remote_peer",
 148                         Network:     rp.Config.ChainID,
 149                         Version:     version.Version,
 150                         ListenAddr:  l.Addr().String(),
 151                         ServiceFlag: consensus.DefaultServices,
 152                 }, 5*time.Second)
 153                 if err != nil {
 154                         fmt.Println("Failed to perform handshake:", err)
 155                 }
 156                 conns = append(conns, conn)
 157                 select {
 158                 case <-rp.quit:
 159                         for _, conn := range conns {
 160                                 if err := conn.Close(); err != nil {
 161                                         fmt.Println(err)
 162                                 }
 163                         }
 164                         return
 165                 default:
 166                 }
 167         }
 168 }
 169
 170 type inboundPeer struct {
 171         PrivKey signlib.PrivKey
 172         config  *cfg.Config
 173 }
 174
 175 func (ip *inboundPeer) dial(addr *NetAddress) {
 176         pc, err := newOutboundPeerConn(addr, ip.PrivKey, DefaultPeerConfig(ip.config.P2P))
 177         if err != nil {
 178                 fmt.Println("newOutboundPeerConn:", err)
 179                 return
 180         }
 181
 182         _, err = pc.HandshakeTimeout(&NodeInfo{
 183                 PubKey:     ip.PrivKey.XPub().String(),
 184                 Moniker:    "remote_peer",
 185                 Network:    ip.config.ChainID,
 186                 Version:    version.Version,
 187                 ListenAddr: addr.String(),
 188         }, 5*time.Second)
 189         if err != nil {
 190                 fmt.Println("Failed to perform handshake:", err)
 191                 return
 192         }
 193         time.AfterFunc(10*time.Second, pc.CloseConn)
 194 }