Discv5 E2E

discv.txt

func createListener(ipAddr net.IP, privKey *ecdsa.PrivateKey, cfg *Config) *discover.UDPv5 {
	udpAddr := &net.UDPAddr{
		IP:   ipAddr,
		Port: int(cfg.Port),
	}
	conn, err := net.ListenUDP("udp4", udpAddr)
	if err != nil {
		log.Fatal(err)
	}
	localNode, err := createLocalNode(privKey, ipAddr, int(cfg.Port))
	if err != nil {
		log.Fatal(err)
	}
	dv5Cfg := discover.Config{
		PrivateKey: privKey,
	}
	if cfg.BootstrapNodeAddr != "" {
		bootNode, err := enode.Parse(enode.ValidSchemes, cfg.BootstrapNodeAddr)
		if err != nil {
			log.Fatal(err)
		}
		dv5Cfg.Bootnodes = []*enode.Node{bootNode}
	}

	network, err := discover.ListenV5(conn, localNode, dv5Cfg)
	if err != nil {
		log.Fatal(err)
	}
	return network
}

func createLocalNode(privKey *ecdsa.PrivateKey, ipAddr net.IP, port int) (*enode.LocalNode, error) {
	db, err := enode.OpenDB("")
	if err != nil {
		return nil, errors.Wrap(err, "Could not open node's peer database")
	}
	localNode := enode.NewLocalNode(db, privKey)
	ipEntry := enr.IP(ipAddr)
	udpEntry := enr.UDP(port)
	localNode.Set(ipEntry)
	localNode.Set(udpEntry)

	return localNode, nil
}

func TestStartDiscV5_DiscoverAllPeers(t *testing.T) {
	port := 2000
	ipAddr, pkey := createAddrAndPrivKey(t)
	bootListener := createListener(ipAddr, pkey, &Config{Port: uint(port)})
	defer bootListener.Close()

	bootNode := bootListener.Self()

	cfg := &Config{
		BootstrapNodeAddr: bootNode.String(),
		Encoding:          "ssz",
	}

	var listeners []*discover.UDPv5
	for i := 1; i <= 1; i++ {
		port = 3000 + i
		cfg.Port = uint(port)
		ipAddr, pkey := createAddrAndPrivKey(t)
		listener, err := createListener(ipAddr, pkey, cfg)
		if err != nil {
			t.Errorf("Could not start discovery for node: %v", err)
		}
		listeners = append(listeners, listener)
	}

	// Wait for the nodes to have their local routing tables to be populated with the other nodes
	time.Sleep(100 * time.Millisecond)

	lastListener := listeners[len(listeners)-1]
	nodes := lastListener.Lookup(bootNode.ID())
	if len(nodes) != 6 {
		t.Errorf("The node's local table doesn't have the expected number of nodes. "+
			"Expected %d but got %d", 6, len(nodes))
	}
	// Close all ports
	for _, listener := range listeners {
		listener.Close()
	}
}