github.com/ethereum/go-ethereum@v1.16.1/cmd/devp2p/internal/v4test/framework.go (about)

     1  // Copyright 2020 The go-ethereum Authors
     2  // This file is part of go-ethereum.
     3  //
     4  // go-ethereum 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-ethereum 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-ethereum. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package v4test
    18  
    19  import (
    20  	"crypto/ecdsa"
    21  	"fmt"
    22  	"net"
    23  	"time"
    24  
    25  	"github.com/ethereum/go-ethereum/crypto"
    26  	"github.com/ethereum/go-ethereum/p2p/discover/v4wire"
    27  	"github.com/ethereum/go-ethereum/p2p/enode"
    28  )
    29  
    30  const waitTime = 300 * time.Millisecond
    31  
    32  type testenv struct {
    33  	l1, l2     net.PacketConn
    34  	key        *ecdsa.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()
    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.IPAddr().IsValid() || node.UDP() == 0 {
    57  		var ip net.IP
    58  		var tcpPort, udpPort int
    59  		if node.IPAddr().IsValid() {
    60  			ip = node.IPAddr().AsSlice()
    61  		} else {
    62  			ip = net.ParseIP("127.0.0.1")
    63  		}
    64  		if tcpPort = node.TCP(); tcpPort == 0 {
    65  			tcpPort = 30303
    66  		}
    67  		if udpPort = node.UDP(); udpPort == 0 {
    68  			udpPort = 30303
    69  		}
    70  		node = enode.NewV4(node.Pubkey(), ip, tcpPort, udpPort)
    71  	}
    72  	addr := &net.UDPAddr{IP: node.IP(), Port: node.UDP()}
    73  	return &testenv{l1, l2, key, node, addr}
    74  }
    75  
    76  func (te *testenv) close() {
    77  	te.l1.Close()
    78  	te.l2.Close()
    79  }
    80  
    81  func (te *testenv) send(c net.PacketConn, req v4wire.Packet) []byte {
    82  	packet, hash, err := v4wire.Encode(te.key, req)
    83  	if err != nil {
    84  		panic(fmt.Errorf("can't encode %v packet: %v", req.Name(), err))
    85  	}
    86  	if _, err := c.WriteTo(packet, te.remoteAddr); err != nil {
    87  		panic(fmt.Errorf("can't send %v: %v", req.Name(), err))
    88  	}
    89  	return hash
    90  }
    91  
    92  func (te *testenv) read(c net.PacketConn) (v4wire.Packet, []byte, error) {
    93  	buf := make([]byte, 2048)
    94  	if err := c.SetReadDeadline(time.Now().Add(waitTime)); err != nil {
    95  		return nil, nil, err
    96  	}
    97  	n, _, err := c.ReadFrom(buf)
    98  	if err != nil {
    99  		return nil, nil, err
   100  	}
   101  	p, _, hash, err := v4wire.Decode(buf[:n])
   102  	return p, hash, err
   103  }
   104  
   105  func (te *testenv) localEndpoint(c net.PacketConn) v4wire.Endpoint {
   106  	addr := c.LocalAddr().(*net.UDPAddr)
   107  	return v4wire.Endpoint{
   108  		IP:  addr.IP.To4(),
   109  		UDP: uint16(addr.Port),
   110  		TCP: 0,
   111  	}
   112  }
   113  
   114  func (te *testenv) remoteEndpoint() v4wire.Endpoint {
   115  	return v4wire.NewEndpoint(te.remoteAddr.AddrPort(), 0)
   116  }
   117  
   118  func contains(ns []v4wire.Node, key v4wire.Pubkey) bool {
   119  	for _, n := range ns {
   120  		if n.ID == key {
   121  			return true
   122  		}
   123  	}
   124  	return false
   125  }