github.com/bytom/bytom@v1.1.2-0.20221014091027-bbcba3df6075/p2p/peer_test.go (about)

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