github.com/inphi/go-ethereum@v1.9.7/p2p/discover/table_util_test.go (about)

     1  // Copyright 2018 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser 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  // The go-ethereum library 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 Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package discover
    18  
    19  import (
    20  	"bytes"
    21  	"crypto/ecdsa"
    22  	"encoding/hex"
    23  	"errors"
    24  	"fmt"
    25  	"math/rand"
    26  	"net"
    27  	"reflect"
    28  	"sort"
    29  	"sync"
    30  
    31  	"github.com/ethereum/go-ethereum/crypto"
    32  	"github.com/ethereum/go-ethereum/log"
    33  	"github.com/ethereum/go-ethereum/p2p/enode"
    34  	"github.com/ethereum/go-ethereum/p2p/enr"
    35  )
    36  
    37  var nullNode *enode.Node
    38  
    39  func init() {
    40  	var r enr.Record
    41  	r.Set(enr.IP{0, 0, 0, 0})
    42  	nullNode = enode.SignNull(&r, enode.ID{})
    43  }
    44  
    45  func newTestTable(t transport) (*Table, *enode.DB) {
    46  	db, _ := enode.OpenDB("")
    47  	tab, _ := newTable(t, db, nil, log.Root())
    48  	go tab.loop()
    49  	return tab, db
    50  }
    51  
    52  // nodeAtDistance creates a node for which enode.LogDist(base, n.id) == ld.
    53  func nodeAtDistance(base enode.ID, ld int, ip net.IP) *node {
    54  	var r enr.Record
    55  	r.Set(enr.IP(ip))
    56  	return wrapNode(enode.SignNull(&r, idAtDistance(base, ld)))
    57  }
    58  
    59  // idAtDistance returns a random hash such that enode.LogDist(a, b) == n
    60  func idAtDistance(a enode.ID, n int) (b enode.ID) {
    61  	if n == 0 {
    62  		return a
    63  	}
    64  	// flip bit at position n, fill the rest with random bits
    65  	b = a
    66  	pos := len(a) - n/8 - 1
    67  	bit := byte(0x01) << (byte(n%8) - 1)
    68  	if bit == 0 {
    69  		pos++
    70  		bit = 0x80
    71  	}
    72  	b[pos] = a[pos]&^bit | ^a[pos]&bit // TODO: randomize end bits
    73  	for i := pos + 1; i < len(a); i++ {
    74  		b[i] = byte(rand.Intn(255))
    75  	}
    76  	return b
    77  }
    78  
    79  func intIP(i int) net.IP {
    80  	return net.IP{byte(i), 0, 2, byte(i)}
    81  }
    82  
    83  // fillBucket inserts nodes into the given bucket until it is full.
    84  func fillBucket(tab *Table, n *node) (last *node) {
    85  	ld := enode.LogDist(tab.self().ID(), n.ID())
    86  	b := tab.bucket(n.ID())
    87  	for len(b.entries) < bucketSize {
    88  		b.entries = append(b.entries, nodeAtDistance(tab.self().ID(), ld, intIP(ld)))
    89  	}
    90  	return b.entries[bucketSize-1]
    91  }
    92  
    93  // fillTable adds nodes the table to the end of their corresponding bucket
    94  // if the bucket is not full. The caller must not hold tab.mutex.
    95  func fillTable(tab *Table, nodes []*node) {
    96  	for _, n := range nodes {
    97  		tab.addSeenNode(n)
    98  	}
    99  }
   100  
   101  type pingRecorder struct {
   102  	mu           sync.Mutex
   103  	dead, pinged map[enode.ID]bool
   104  	records      map[enode.ID]*enode.Node
   105  	n            *enode.Node
   106  }
   107  
   108  func newPingRecorder() *pingRecorder {
   109  	var r enr.Record
   110  	r.Set(enr.IP{0, 0, 0, 0})
   111  	n := enode.SignNull(&r, enode.ID{})
   112  
   113  	return &pingRecorder{
   114  		dead:    make(map[enode.ID]bool),
   115  		pinged:  make(map[enode.ID]bool),
   116  		records: make(map[enode.ID]*enode.Node),
   117  		n:       n,
   118  	}
   119  }
   120  
   121  // setRecord updates a node record. Future calls to ping and
   122  // requestENR will return this record.
   123  func (t *pingRecorder) updateRecord(n *enode.Node) {
   124  	t.mu.Lock()
   125  	defer t.mu.Unlock()
   126  	t.records[n.ID()] = n
   127  }
   128  
   129  // Stubs to satisfy the transport interface.
   130  func (t *pingRecorder) Self() *enode.Node           { return nullNode }
   131  func (t *pingRecorder) lookupSelf() []*enode.Node   { return nil }
   132  func (t *pingRecorder) lookupRandom() []*enode.Node { return nil }
   133  func (t *pingRecorder) close()                      {}
   134  
   135  // ping simulates a ping request.
   136  func (t *pingRecorder) ping(n *enode.Node) (seq uint64, err error) {
   137  	t.mu.Lock()
   138  	defer t.mu.Unlock()
   139  
   140  	t.pinged[n.ID()] = true
   141  	if t.dead[n.ID()] {
   142  		return 0, errTimeout
   143  	}
   144  	if t.records[n.ID()] != nil {
   145  		seq = t.records[n.ID()].Seq()
   146  	}
   147  	return seq, nil
   148  }
   149  
   150  // requestENR simulates an ENR request.
   151  func (t *pingRecorder) RequestENR(n *enode.Node) (*enode.Node, error) {
   152  	t.mu.Lock()
   153  	defer t.mu.Unlock()
   154  
   155  	if t.dead[n.ID()] || t.records[n.ID()] == nil {
   156  		return nil, errTimeout
   157  	}
   158  	return t.records[n.ID()], nil
   159  }
   160  
   161  func hasDuplicates(slice []*node) bool {
   162  	seen := make(map[enode.ID]bool)
   163  	for i, e := range slice {
   164  		if e == nil {
   165  			panic(fmt.Sprintf("nil *Node at %d", i))
   166  		}
   167  		if seen[e.ID()] {
   168  			return true
   169  		}
   170  		seen[e.ID()] = true
   171  	}
   172  	return false
   173  }
   174  
   175  func checkNodesEqual(got, want []*enode.Node) error {
   176  	if reflect.DeepEqual(got, want) {
   177  		return nil
   178  	}
   179  	output := new(bytes.Buffer)
   180  	fmt.Fprintf(output, "got %d nodes:\n", len(got))
   181  	for _, n := range got {
   182  		fmt.Fprintf(output, "  %v %v\n", n.ID(), n)
   183  	}
   184  	fmt.Fprintf(output, "want %d:\n", len(want))
   185  	for _, n := range want {
   186  		fmt.Fprintf(output, "  %v %v\n", n.ID(), n)
   187  	}
   188  	return errors.New(output.String())
   189  }
   190  
   191  func sortByID(nodes []*enode.Node) {
   192  	sort.Slice(nodes, func(i, j int) bool {
   193  		return string(nodes[i].ID().Bytes()) < string(nodes[j].ID().Bytes())
   194  	})
   195  }
   196  
   197  func sortedByDistanceTo(distbase enode.ID, slice []*node) bool {
   198  	return sort.SliceIsSorted(slice, func(i, j int) bool {
   199  		return enode.DistCmp(distbase, slice[i].ID(), slice[j].ID()) < 0
   200  	})
   201  }
   202  
   203  func hexEncPrivkey(h string) *ecdsa.PrivateKey {
   204  	b, err := hex.DecodeString(h)
   205  	if err != nil {
   206  		panic(err)
   207  	}
   208  	key, err := crypto.ToECDSA(b)
   209  	if err != nil {
   210  		panic(err)
   211  	}
   212  	return key
   213  }
   214  
   215  func hexEncPubkey(h string) (ret encPubkey) {
   216  	b, err := hex.DecodeString(h)
   217  	if err != nil {
   218  		panic(err)
   219  	}
   220  	if len(b) != len(ret) {
   221  		panic("invalid length")
   222  	}
   223  	copy(ret[:], b)
   224  	return ret
   225  }