github.com/core-coin/go-core/v2@v2.1.9/cmd/devp2p/internal/v4test/framework.go (about)

     1  // Copyright 2020 by the Authors
     2  // This file is part of go-core.
     3  //
     4  // go-core is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // go-core is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU General Public License
    15  // along with go-core. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package v4test
    18  
    19  import (
    20  	crand "crypto/rand"
    21  	"fmt"
    22  	"net"
    23  	"time"
    24  
    25  	"github.com/core-coin/go-core/v2/crypto"
    26  	"github.com/core-coin/go-core/v2/p2p/discover/v4wire"
    27  	"github.com/core-coin/go-core/v2/p2p/enode"
    28  )
    29  
    30  const waitTime = 300 * time.Millisecond
    31  
    32  type testenv struct {
    33  	l1, l2     net.PacketConn
    34  	key        *crypto.PrivateKey
    35  	remote     *enode.Node
    36  	remoteAddr *net.UDPAddr
    37  }
    38  
    39  func newTestEnv(remote string, listen1, listen2 string) *testenv {
    40  	l1, err := net.ListenPacket("udp", fmt.Sprintf("%v:0", listen1))
    41  	if err != nil {
    42  		panic(err)
    43  	}
    44  	l2, err := net.ListenPacket("udp", fmt.Sprintf("%v:0", listen2))
    45  	if err != nil {
    46  		panic(err)
    47  	}
    48  	key, err := crypto.GenerateKey(crand.Reader)
    49  	if err != nil {
    50  		panic(err)
    51  	}
    52  	node, err := enode.Parse(enode.ValidSchemes, remote)
    53  	if err != nil {
    54  		panic(err)
    55  	}
    56  	if node.IP() == nil || node.UDP() == 0 {
    57  		var ip net.IP
    58  		var tcpPort, udpPort int
    59  		if ip = node.IP(); ip == nil {
    60  			ip = net.ParseIP("127.0.0.1")
    61  		}
    62  		if tcpPort = node.TCP(); tcpPort == 0 {
    63  			tcpPort = 30300
    64  		}
    65  		if udpPort = node.TCP(); udpPort == 0 {
    66  			udpPort = 30300
    67  		}
    68  		node = enode.NewV4(node.Pubkey(), ip, tcpPort, udpPort)
    69  	}
    70  	addr := &net.UDPAddr{IP: node.IP(), Port: node.UDP()}
    71  	return &testenv{l1, l2, key, node, addr}
    72  }
    73  
    74  func (te *testenv) close() {
    75  	te.l1.Close()
    76  	te.l2.Close()
    77  }
    78  
    79  func (te *testenv) send(c net.PacketConn, req v4wire.Packet) []byte {
    80  	packet, hash, err := v4wire.Encode(te.key, req)
    81  	if err != nil {
    82  		panic(fmt.Errorf("can't encode %v packet: %v", req.Name(), err))
    83  	}
    84  	if _, err := c.WriteTo(packet, te.remoteAddr); err != nil {
    85  		panic(fmt.Errorf("can't send %v: %v", req.Name(), err))
    86  	}
    87  	return hash
    88  }
    89  
    90  func (te *testenv) read(c net.PacketConn) (v4wire.Packet, []byte, error) {
    91  	buf := make([]byte, 2048)
    92  	if err := c.SetReadDeadline(time.Now().Add(waitTime)); err != nil {
    93  		return nil, nil, err
    94  	}
    95  	n, _, err := c.ReadFrom(buf)
    96  	if err != nil {
    97  		return nil, nil, err
    98  	}
    99  	p, _, hash, err := v4wire.Decode(buf[:n])
   100  	return p, hash, err
   101  }
   102  
   103  func (te *testenv) localEndpoint(c net.PacketConn) v4wire.Endpoint {
   104  	addr := c.LocalAddr().(*net.UDPAddr)
   105  	return v4wire.Endpoint{
   106  		IP:  addr.IP.To4(),
   107  		UDP: uint16(addr.Port),
   108  		TCP: 0,
   109  	}
   110  }
   111  
   112  func (te *testenv) remoteEndpoint() v4wire.Endpoint {
   113  	return v4wire.NewEndpoint(te.remoteAddr, 0)
   114  }
   115  
   116  func contains(ns []v4wire.Node, key v4wire.Pubkey) bool {
   117  	for _, n := range ns {
   118  		if n.ID == key {
   119  			return true
   120  		}
   121  	}
   122  	return false
   123  }