github.com/jonasnick/go-ethereum@v0.7.12-0.20150216215225-22176f05d387/p2p/server_test.go (about)

     1  package p2p
     2  
     3  import (
     4  	"bytes"
     5  	"crypto/ecdsa"
     6  	"io"
     7  	"math/rand"
     8  	"net"
     9  	"sync"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/jonasnick/go-ethereum/crypto"
    14  	"github.com/jonasnick/go-ethereum/p2p/discover"
    15  )
    16  
    17  func startTestServer(t *testing.T, pf newPeerHook) *Server {
    18  	server := &Server{
    19  		Name:        "test",
    20  		MaxPeers:    10,
    21  		ListenAddr:  "127.0.0.1:0",
    22  		PrivateKey:  newkey(),
    23  		newPeerHook: pf,
    24  		handshakeFunc: func(io.ReadWriter, *ecdsa.PrivateKey, *discover.Node) (id discover.NodeID, st []byte, err error) {
    25  			return randomID(), nil, err
    26  		},
    27  	}
    28  	if err := server.Start(); err != nil {
    29  		t.Fatalf("Could not start server: %v", err)
    30  	}
    31  	return server
    32  }
    33  
    34  func TestServerListen(t *testing.T) {
    35  	defer testlog(t).detach()
    36  
    37  	// start the test server
    38  	connected := make(chan *Peer)
    39  	srv := startTestServer(t, func(p *Peer) {
    40  		if p == nil {
    41  			t.Error("peer func called with nil conn")
    42  		}
    43  		connected <- p
    44  	})
    45  	defer close(connected)
    46  	defer srv.Stop()
    47  
    48  	// dial the test server
    49  	conn, err := net.DialTimeout("tcp", srv.ListenAddr, 5*time.Second)
    50  	if err != nil {
    51  		t.Fatalf("could not dial: %v", err)
    52  	}
    53  	defer conn.Close()
    54  
    55  	select {
    56  	case peer := <-connected:
    57  		if peer.LocalAddr().String() != conn.RemoteAddr().String() {
    58  			t.Errorf("peer started with wrong conn: got %v, want %v",
    59  				peer.LocalAddr(), conn.RemoteAddr())
    60  		}
    61  	case <-time.After(1 * time.Second):
    62  		t.Error("server did not accept within one second")
    63  	}
    64  }
    65  
    66  func TestServerDial(t *testing.T) {
    67  	defer testlog(t).detach()
    68  
    69  	// run a one-shot TCP server to handle the connection.
    70  	listener, err := net.Listen("tcp", "127.0.0.1:0")
    71  	if err != nil {
    72  		t.Fatalf("could not setup listener: %v")
    73  	}
    74  	defer listener.Close()
    75  	accepted := make(chan net.Conn)
    76  	go func() {
    77  		conn, err := listener.Accept()
    78  		if err != nil {
    79  			t.Error("accept error:", err)
    80  			return
    81  		}
    82  		conn.Close()
    83  		accepted <- conn
    84  	}()
    85  
    86  	// start the server
    87  	connected := make(chan *Peer)
    88  	srv := startTestServer(t, func(p *Peer) { connected <- p })
    89  	defer close(connected)
    90  	defer srv.Stop()
    91  
    92  	// tell the server to connect
    93  	tcpAddr := listener.Addr().(*net.TCPAddr)
    94  	srv.SuggestPeer(&discover.Node{IP: tcpAddr.IP, TCPPort: tcpAddr.Port})
    95  
    96  	select {
    97  	case conn := <-accepted:
    98  		select {
    99  		case peer := <-connected:
   100  			if peer.RemoteAddr().String() != conn.LocalAddr().String() {
   101  				t.Errorf("peer started with wrong conn: got %v, want %v",
   102  					peer.RemoteAddr(), conn.LocalAddr())
   103  			}
   104  			// TODO: validate more fields
   105  		case <-time.After(1 * time.Second):
   106  			t.Error("server did not launch peer within one second")
   107  		}
   108  
   109  	case <-time.After(1 * time.Second):
   110  		t.Error("server did not connect within one second")
   111  	}
   112  }
   113  
   114  func TestServerBroadcast(t *testing.T) {
   115  	defer testlog(t).detach()
   116  
   117  	var connected sync.WaitGroup
   118  	srv := startTestServer(t, func(p *Peer) {
   119  		p.protocols = []Protocol{discard}
   120  		p.startSubprotocols([]Cap{discard.cap()})
   121  		p.noHandshake = true
   122  		connected.Done()
   123  	})
   124  	defer srv.Stop()
   125  
   126  	// create a few peers
   127  	var conns = make([]net.Conn, 8)
   128  	connected.Add(len(conns))
   129  	deadline := time.Now().Add(3 * time.Second)
   130  	dialer := &net.Dialer{Deadline: deadline}
   131  	for i := range conns {
   132  		conn, err := dialer.Dial("tcp", srv.ListenAddr)
   133  		if err != nil {
   134  			t.Fatalf("conn %d: dial error: %v", i, err)
   135  		}
   136  		defer conn.Close()
   137  		conn.SetDeadline(deadline)
   138  		conns[i] = conn
   139  	}
   140  	connected.Wait()
   141  
   142  	// broadcast one message
   143  	srv.Broadcast("discard", 0, "foo")
   144  	goldbuf := new(bytes.Buffer)
   145  	writeMsg(goldbuf, NewMsg(16, "foo"))
   146  	golden := goldbuf.Bytes()
   147  
   148  	// check that the message has been written everywhere
   149  	for i, conn := range conns {
   150  		buf := make([]byte, len(golden))
   151  		if _, err := io.ReadFull(conn, buf); err != nil {
   152  			t.Errorf("conn %d: read error: %v", i, err)
   153  		} else if !bytes.Equal(buf, golden) {
   154  			t.Errorf("conn %d: msg mismatch\ngot:  %x\nwant: %x", i, buf, golden)
   155  		}
   156  	}
   157  }
   158  
   159  func newkey() *ecdsa.PrivateKey {
   160  	key, err := crypto.GenerateKey()
   161  	if err != nil {
   162  		panic("couldn't generate key: " + err.Error())
   163  	}
   164  	return key
   165  }
   166  
   167  func randomID() (id discover.NodeID) {
   168  	for i := range id {
   169  		id[i] = byte(rand.Intn(255))
   170  	}
   171  	return id
   172  }