github.com/turingchain2020/turingchain@v1.1.21/system/p2p/dht/extension/relay_test.go (about)

     1  package extension
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"io"
     7  	"net"
     8  	"testing"
     9  	"time"
    10  
    11  	bhost "github.com/libp2p/go-libp2p-blankhost"
    12  	circuit "github.com/libp2p/go-libp2p-circuit"
    13  	"github.com/libp2p/go-libp2p-core/host"
    14  	discovery "github.com/libp2p/go-libp2p-discovery"
    15  	dht "github.com/libp2p/go-libp2p-kad-dht"
    16  	swarmt "github.com/libp2p/go-libp2p-swarm/testing"
    17  	"github.com/stretchr/testify/require"
    18  )
    19  
    20  func getNetHosts(ctx context.Context, n int, t *testing.T) []host.Host {
    21  	var out []host.Host
    22  
    23  	for i := 0; i < n; i++ {
    24  		netw := swarmt.GenSwarm(t, ctx)
    25  		h := bhost.NewBlankHost(netw)
    26  		out = append(out, h)
    27  	}
    28  
    29  	return out
    30  }
    31  
    32  func connect(t *testing.T, a, b host.Host) {
    33  	pinfo := a.Peerstore().PeerInfo(a.ID())
    34  	err := b.Connect(context.Background(), pinfo)
    35  	if err != nil {
    36  		t.Fatal(err)
    37  	}
    38  }
    39  
    40  func TestRelay(t *testing.T) {
    41  	ctx, cancel := context.WithCancel(context.Background())
    42  	defer cancel()
    43  
    44  	hosts := getNetHosts(ctx, 3, t)
    45  
    46  	connect(t, hosts[0], hosts[1])
    47  	connect(t, hosts[1], hosts[2])
    48  	//第二个节点作为中继节点
    49  	newRelay(ctx, hosts[1], circuit.OptHop)
    50  	r2, err := newRelay(ctx, hosts[2])
    51  	require.Nil(t, err)
    52  
    53  	var (
    54  		conn1, conn2 net.Conn
    55  		done         = make(chan struct{})
    56  	)
    57  
    58  	defer func() {
    59  		<-done
    60  		if conn1 != nil {
    61  			conn1.Close()
    62  		}
    63  		if conn2 != nil {
    64  			conn2.Close()
    65  		}
    66  	}()
    67  	msg := []byte("relay works!")
    68  	go func() {
    69  		defer close(done)
    70  		//第三个节点监听
    71  		list := r2.Listener()
    72  
    73  		var err error
    74  		conn1, err = list.Accept()
    75  		if err != nil {
    76  			t.Error(err)
    77  			return
    78  		}
    79  
    80  		_, err = conn1.Write(msg)
    81  		if err != nil {
    82  			t.Error(err)
    83  			return
    84  		}
    85  	}()
    86  
    87  	rinfo := hosts[1].Peerstore().PeerInfo(hosts[1].ID()) //中继节点的peerinfo
    88  	require.NotNil(t, rinfo.Addrs)
    89  	dinfo := hosts[2].Peerstore().PeerInfo(hosts[2].ID()) //目的节点
    90  	require.NotNil(t, dinfo.Addrs)
    91  	kademliaDHT, err := dht.New(context.Background(), hosts[0])
    92  	require.Nil(t, err)
    93  	_, err = kademliaDHT.RoutingTable().Update(hosts[1].ID())
    94  	require.Nil(t, err)
    95  	netRely := NewRelayDiscovery(hosts[0], discovery.NewRoutingDiscovery(kademliaDHT))
    96  	netRely.Advertise(ctx)
    97  	conn2, err = netRely.DialDestPeer(rinfo, dinfo)
    98  	if err != nil {
    99  		t.Fatal(err)
   100  	}
   101  
   102  	err = conn2.SetReadDeadline(time.Now().Add(time.Second))
   103  	require.Nil(t, err)
   104  	result := make([]byte, len(msg))
   105  	_, err = io.ReadFull(conn2, result)
   106  	if err != nil {
   107  		t.Fatal(err)
   108  	}
   109  
   110  	if !bytes.Equal(result, msg) {
   111  		t.Fatal("message was incorrect:", string(result))
   112  	}
   113  
   114  	testCheckOp(t, netRely, hosts[1])
   115  	testFindOpPeers(t, netRely)
   116  
   117  }
   118  
   119  func testCheckOp(t *testing.T, netRely *Relay, h host.Host) {
   120  	//check op
   121  	ok, err := netRely.CheckHOp(h.ID())
   122  	require.Nil(t, err)
   123  	require.True(t, ok)
   124  }
   125  
   126  func testFindOpPeers(t *testing.T, netRely *Relay) {
   127  	peers, err := netRely.FindOpPeers()
   128  	if err != nil {
   129  		t.Fatal(err)
   130  	}
   131  	t.Log(peers)
   132  }