package p2p
import (
"crypto/ecdsa"
"math/rand"
"net"
"testing"
"time"
"github.com/ethereum/go-ethereum/crypto"
"github.com/ethereum/go-ethereum/crypto/sha3"
"github.com/ethereum/go-ethereum/p2p/discover"
)
func startTestServer(t *testing.T, pf newPeerHook) *Server {
server := &Server{
Name: "test",
MaxPeers: 10,
ListenAddr: "127.0.0.1:0",
PrivateKey: newkey(),
newPeerHook: pf,
setupFunc: func(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, dial *discover.Node, keepconn func(discover.NodeID) bool) (*conn, error) {
id := randomID()
if !keepconn(id) {
return nil, DiscAlreadyConnected
}
rw := newRlpxFrameRW(fd, secrets{
MAC: zero16,
AES: zero16,
IngressMAC: sha3.NewKeccak256(),
EgressMAC: sha3.NewKeccak256(),
})
return &conn{
MsgReadWriter: rw,
protoHandshake: &protoHandshake{ID: id, Version: baseProtocolVersion},
}, nil
},
}
if err := server.Start(); err != nil {
t.Fatalf("Could not start server: %v", err)
}
return server
}
func TestServerListen(t *testing.T) {
// start the test server
connected := make(chan *Peer)
srv := startTestServer(t, func(p *Peer) {
if p == nil {
t.Error("peer func called with nil conn")
}
connected <- p
})
defer close(connected)
defer srv.Stop()
// dial the test server
conn, err := net.DialTimeout("tcp", srv.ListenAddr, 5*time.Second)
if err != nil {
t.Fatalf("could not dial: %v", err)
}
defer conn.Close()
select {
case peer := <-connected:
if peer.LocalAddr().String() != conn.RemoteAddr().String() {
t.Errorf("peer started with wrong conn: got %v, want %v",
peer.LocalAddr(), conn.RemoteAddr())
}
case <-time.After(1 * time.Second):
t.Error("server did not accept within one second")
}
}
func TestServerDial(t *testing.T) {
// run a one-shot TCP server to handle the connection.
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("could not setup listener: %v")
}
defer listener.Close()
accepted := make(chan net.Conn)
go func() {
conn, err := listener.Accept()
if err != nil {
t.Error("accept error:", err)
return
}
conn.Close()
accepted <- conn
}()
// start the server
connected := make(chan *Peer)
srv := startTestServer(t, func(p *Peer) { connected <- p })
defer close(connected)
defer srv.Stop()
// tell the server to connect
tcpAddr := listener.Addr().(*net.TCPAddr)
srv.staticDial <- &discover.Node{IP: tcpAddr.IP, TCP: uint16(tcpAddr.Port)}
select {
case conn := <-accepted:
select {
case peer := <-connected:
if peer.RemoteAddr().String() != conn.LocalAddr().String() {
t.Errorf("peer started with wrong conn: got %v, want %v",
peer.RemoteAddr(), conn.LocalAddr())
}
// TODO: validate more fields
case <-time.After(1 * time.Second):
t.Error("server did not launch peer within one second")
}
case <-time.After(1 * time.Second):
t.Error("server did not connect within one second")
}
}
// This test checks that connections are disconnected
// just after the encryption handshake when the server is
// at capacity.
//
// It also serves as a light-weight integration test.
func TestServerDisconnectAtCap(t *testing.T) {
started := make(chan *Peer)
srv := &Server{
ListenAddr: "127.0.0.1:0",
PrivateKey: newkey(),
MaxPeers: 10,
NoDial: true,
// This hook signals that the peer was actually started. We
// need to wait for the peer to be started before dialing the
// next connection to get a deterministic peer count.
newPeerHook: func(p *Peer) { started <- p },
}
if err := srv.Start(); err != nil {
t.Fatal(err)
}
defer srv.Stop()
nconns := srv.MaxPeers + 1
dialer := &net.Dialer{Deadline: time.Now().Add(3 * time.Second)}
for i := 0; i < nconns; i++ {
conn, err := dialer.Dial("tcp", srv.ListenAddr)
if err != nil {
t.Fatalf("conn %d: dial error: %v", i, err)
}
// Close the connection when the test ends, before
// shutting down the server.
defer conn.Close()
// Run the handshakes just like a real peer would.
key := newkey()
hs := &protoHandshake{Version: baseProtocolVersion, ID: discover.PubkeyID(&key.PublicKey)}
_, err = setupConn(conn, key, hs, srv.Self(), keepalways)
if i == nconns-1 {
// When handling the last connection, the server should
// disconnect immediately instead of running the protocol
// handshake.
if err != DiscTooManyPeers {
t.Errorf("conn %d: got error %q, expected %q", i, err, DiscTooManyPeers)
}
} else {
// For all earlier connections, the handshake should go through.
if err != nil {
t.Fatalf("conn %d: unexpected error: %v", i, err)
}
// Wait for runPeer to be started.
<-started
}
}
}
// Tests that static peers are (re)connected, and done so even above max peers.
func TestServerStaticPeers(t *testing.T) {
// Create a test server with limited connection slots
started := make(chan *Peer)
server := &Server{
ListenAddr: "127.0.0.1:0",
PrivateKey: newkey(),
MaxPeers: 3,
newPeerHook: func(p *Peer) { started <- p },
staticCycle: time.Second,
}
if err := server.Start(); err != nil {
t.Fatal(err)
}
defer server.Stop()
// Fill up all the slots on the server
dialer := &net.Dialer{Deadline: time.Now().Add(3 * time.Second)}
for i := 0; i < server.MaxPeers; i++ {
// Establish a new connection
conn, err := dialer.Dial("tcp", server.ListenAddr)
if err != nil {
t.Fatalf("conn %d: dial error: %v", i, err)
}
defer conn.Close()
// Run the handshakes just like a real peer would, and wait for completion
key := newkey()
shake := &protoHandshake{Version: baseProtocolVersion, ID: discover.PubkeyID(&key.PublicKey)}
if _, err = setupConn(conn, key, shake, server.Self(), keepalways); err != nil {
t.Fatalf("conn %d: unexpected error: %v", i, err)
}
<-started
}
// Open a TCP listener to accept static connections
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("failed to setup listener: %v", err)
}
defer listener.Close()
connected := make(chan net.Conn)
go func() {
for i := 0; i < 3; i++ {
conn, err := listener.Accept()
if err == nil {
connected <- conn
}
}
}()
// Inject a static node and wait for a remote dial, then redial, then nothing
addr := listener.Addr().(*net.TCPAddr)
static := &discover.Node{
ID: discover.PubkeyID(&newkey().PublicKey),
IP: addr.IP,
TCP: uint16(addr.Port),
}
server.AddPeer(static)
select {
case conn := <-connected:
// Close the first connection, expect redial
conn.Close()
case <-time.After(2 * server.staticCycle):
t.Fatalf("remote dial timeout")
}
select {
case conn := <-connected:
// Keep the second connection, don't expect redial
defer conn.Close()
case <-time.After(2 * server.staticCycle):
t.Fatalf("remote re-dial timeout")
}
select {
case <-time.After(2 * server.staticCycle):
// Timeout as no dial occurred
case <-connected:
t.Fatalf("connected node dialed")
}
}
// Tests that trusted peers and can connect above max peer caps.
func TestServerTrustedPeers(t *testing.T) {
// Create a trusted peer to accept connections from
key := newkey()
trusted := &discover.Node{
ID: discover.PubkeyID(&key.PublicKey),
}
// Create a test server with limited connection slots
started := make(chan *Peer)
server := &Server{
ListenAddr: "127.0.0.1:0",
PrivateKey: newkey(),
MaxPeers: 3,
NoDial: true,
TrustedNodes: []*discover.Node{trusted},
newPeerHook: func(p *Peer) { started <- p },
}
if err := server.Start(); err != nil {
t.Fatal(err)
}
defer server.Stop()
// Fill up all the slots on the server
dialer := &net.Dialer{Deadline: time.Now().Add(3 * time.Second)}
for i := 0; i < server.MaxPeers; i++ {
// Establish a new connection
conn, err := dialer.Dial("tcp", server.ListenAddr)
if err != nil {
t.Fatalf("conn %d: dial error: %v", i, err)
}
defer conn.Close()
// Run the handshakes just like a real peer would, and wait for completion
key := newkey()
shake := &protoHandshake{Version: baseProtocolVersion, ID: discover.PubkeyID(&key.PublicKey)}
if _, err = setupConn(conn, key, shake, server.Self(), keepalways); err != nil {
t.Fatalf("conn %d: unexpected error: %v", i, err)
}
<-started
}
// Dial from the trusted peer, ensure connection is accepted
conn, err := dialer.Dial("tcp", server.ListenAddr)
if err != nil {
t.Fatalf("trusted node: dial error: %v", err)
}
defer conn.Close()
shake := &protoHandshake{Version: baseProtocolVersion, ID: trusted.ID}
if _, err = setupConn(conn, key, shake, server.Self(), keepalways); err != nil {
t.Fatalf("trusted node: unexpected error: %v", err)
}
select {
case <-started:
// Ok, trusted peer accepted
case <-time.After(100 * time.Millisecond):
t.Fatalf("trusted node timeout")
}
}
// Tests that a failed dial will temporarily throttle a peer.
func TestServerMaxPendingDials(t *testing.T) {
// Start a simple test server
server := &Server{
ListenAddr: "127.0.0.1:0",
PrivateKey: newkey(),
MaxPeers: 10,
MaxPendingPeers: 1,
}
if err := server.Start(); err != nil {
t.Fatal("failed to start test server: %v", err)
}
defer server.Stop()
// Simulate two separate remote peers
peers := make(chan *discover.Node, 2)
conns := make(chan net.Conn, 2)
for i := 0; i < 2; i++ {
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listener %d: failed to setup: %v", i, err)
}
defer listener.Close()
addr := listener.Addr().(*net.TCPAddr)
peers <- &discover.Node{
ID: discover.PubkeyID(&newkey().PublicKey),
IP: addr.IP,
TCP: uint16(addr.Port),
}
go func() {
conn, err := listener.Accept()
if err == nil {
conns <- conn
}
}()
}
// Request a dial for both peers
go func() {
for i := 0; i < 2; i++ {
server.staticDial <- <-peers // hack piggybacking the static implementation
}
}()
// Make sure only one outbound connection goes through
var conn net.Conn
select {
case conn = <-conns:
case <-time.After(100 * time.Millisecond):
t.Fatalf("first dial timeout")
}
select {
case conn = <-conns:
t.Fatalf("second dial completed prematurely")
case <-time.After(100 * time.Millisecond):
}
// Finish the first dial, check the second
conn.Close()
select {
case conn = <-conns:
conn.Close()
case <-time.After(100 * time.Millisecond):
t.Fatalf("second dial timeout")
}
}
func TestServerMaxPendingAccepts(t *testing.T) {
// Start a test server and a peer sink for synchronization
started := make(chan *Peer)
server := &Server{
ListenAddr: "127.0.0.1:0",
PrivateKey: newkey(),
MaxPeers: 10,
MaxPendingPeers: 1,
NoDial: true,
newPeerHook: func(p *Peer) { started <- p },
}
if err := server.Start(); err != nil {
t.Fatal("failed to start test server: %v", err)
}
defer server.Stop()
// Try and connect to the server on multiple threads concurrently
conns := make([]net.Conn, 2)
for i := 0; i < 2; i++ {
dialer := &net.Dialer{Deadline: time.Now().Add(3 * time.Second)}
conn, err := dialer.Dial("tcp", server.ListenAddr)
if err != nil {
t.Fatalf("failed to dial server: %v", err)
}
conns[i] = conn
}
// Check that a handshake on the second doesn't pass
go func() {
key := newkey()
shake := &protoHandshake{Version: baseProtocolVersion, ID: discover.PubkeyID(&key.PublicKey)}
if _, err := setupConn(conns[1], key, shake, server.Self(), keepalways); err != nil {
t.Fatalf("failed to run handshake: %v", err)
}
}()
select {
case <-started:
t.Fatalf("handshake on second connection accepted")
case <-time.After(time.Second):
}
// Shake on first, check that both go through
go func() {
key := newkey()
shake := &protoHandshake{Version: baseProtocolVersion, ID: discover.PubkeyID(&key.PublicKey)}
if _, err := setupConn(conns[0], key, shake, server.Self(), keepalways); err != nil {
t.Fatalf("failed to run handshake: %v", err)
}
}()
for i := 0; i < 2; i++ {
select {
case <-started:
case <-time.After(time.Second):
t.Fatalf("peer %d: handshake timeout", i)
}
}
}
func newkey() *ecdsa.PrivateKey {
key, err := crypto.GenerateKey()
if err != nil {
panic("couldn't generate key: " + err.Error())
}
return key
}
func randomID() (id discover.NodeID) {
for i := range id {
id[i] = byte(rand.Intn(255))
}
return id
}
func keepalways(id discover.NodeID) bool {
return true
}