github.com/annchain/OG@v0.0.9/p2p/peer_test.go (about)

     1  // Copyright 2014 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 p2p
    18  
    19  import (
    20  	"errors"
    21  	"fmt"
    22  	"github.com/annchain/OG/types/msg"
    23  	"github.com/sirupsen/logrus"
    24  	"math/rand"
    25  	"net"
    26  	"os"
    27  	"reflect"
    28  	"testing"
    29  	"time"
    30  )
    31  
    32  var discard = Protocol{
    33  	Name:   "discard",
    34  	Length: 1,
    35  	Run: func(p *Peer, rw MsgReadWriter) error {
    36  		for {
    37  			msg, err := rw.ReadMsg()
    38  			if err != nil {
    39  				return err
    40  			}
    41  			fmt.Printf("discarding %d\n", msg.Code)
    42  			if err = msg.Discard(); err != nil {
    43  				return err
    44  			}
    45  		}
    46  	},
    47  }
    48  
    49  func testPeer(protos []Protocol) (func(), *conn, *Peer, <-chan error) {
    50  	fd1, fd2 := net.Pipe()
    51  	c1 := &conn{fd: fd1, node: newNode(randomID(), nil), transport: newTestTransport(&newkey().PublicKey, fd1)}
    52  	c2 := &conn{fd: fd2, node: newNode(randomID(), nil), transport: newTestTransport(&newkey().PublicKey, fd2)}
    53  	for _, p := range protos {
    54  		c1.caps = append(c1.caps, p.cap())
    55  		c2.caps = append(c2.caps, p.cap())
    56  	}
    57  
    58  	peer := newPeer(c1, protos)
    59  	errc := make(chan error, 1)
    60  	go func() {
    61  		_, err := peer.run()
    62  		errc <- err
    63  	}()
    64  
    65  	closer := func() { c2.close(errors.New("close func called")) }
    66  	return closer, c2, peer, errc
    67  }
    68  
    69  func TestPeerProtoReadMsgRlp(t *testing.T) {
    70  	proto := Protocol{
    71  		Name:   "a",
    72  		Length: 5,
    73  		Run: func(peer *Peer, rw MsgReadWriter) error {
    74  			if err := ExpectMsg(rw, 2, msg.Uints{1}); err != nil {
    75  				t.Error(err)
    76  			}
    77  			if err := ExpectMsg(rw, 3, msg.Uints{2}); err != nil {
    78  				t.Error(err)
    79  			}
    80  			if err := ExpectMsg(rw, 4, msg.Uints{3}); err != nil {
    81  				t.Error(err)
    82  			}
    83  			return nil
    84  		},
    85  	}
    86  
    87  	closer, rw, _, errc := testPeer([]Protocol{proto})
    88  	defer closer()
    89  	a := msg.Uints{1}
    90  	b, _ := a.MarshalMsg(nil)
    91  	Send(rw, baseProtocolLength+2, b)
    92  	a = msg.Uints{2}
    93  	b, _ = a.MarshalMsg(nil)
    94  	Send(rw, baseProtocolLength+3, b)
    95  	a = msg.Uints{3}
    96  	b, _ = a.MarshalMsg(nil)
    97  	Send(rw, baseProtocolLength+4, b)
    98  
    99  	select {
   100  	case err := <-errc:
   101  		if err != errProtocolReturned {
   102  			t.Errorf("peer returned error: %v", err)
   103  		}
   104  	case <-time.After(2 * time.Second):
   105  		t.Errorf("receive timeout")
   106  	}
   107  }
   108  
   109  func TestPeerProtoReadMsg(t *testing.T) {
   110  	proto := Protocol{
   111  		Name:   "a",
   112  		Length: 5,
   113  		Run: func(peer *Peer, rw MsgReadWriter) error {
   114  			if err := ExpectMsg(rw, 2, msg.Bytes{1}); err != nil {
   115  				t.Error(err)
   116  			}
   117  			if err := ExpectMsg(rw, 3, msg.Bytes{2}); err != nil {
   118  				t.Error(err)
   119  			}
   120  			if err := ExpectMsg(rw, 4, msg.Bytes{3}); err != nil {
   121  				t.Error(err)
   122  			}
   123  			return nil
   124  		},
   125  	}
   126  
   127  	closer, rw, _, errc := testPeer([]Protocol{proto})
   128  	defer closer()
   129  	a := msg.Bytes{1}
   130  	b, _ := a.MarshalMsg(nil)
   131  	Send(rw, baseProtocolLength+2, b)
   132  	a = msg.Bytes{2}
   133  	b, _ = a.MarshalMsg(nil)
   134  	Send(rw, baseProtocolLength+3, b)
   135  	a = msg.Bytes{3}
   136  	b, _ = a.MarshalMsg(nil)
   137  	Send(rw, baseProtocolLength+4, b)
   138  
   139  	select {
   140  	case err := <-errc:
   141  		if err != errProtocolReturned {
   142  			t.Errorf("peer returned error: %v", err)
   143  		}
   144  	case <-time.After(2 * time.Second):
   145  		t.Errorf("receive timeout")
   146  	}
   147  }
   148  func TestPeerProtoEncodeMsg(t *testing.T) {
   149  	proto := Protocol{
   150  		Name:   "a",
   151  		Length: 2,
   152  		Run: func(peer *Peer, rw MsgReadWriter) error {
   153  			if err := Send(rw, 2, nil); err == nil {
   154  				t.Error("expected error for out-of-range msg code, got nil")
   155  			}
   156  			a := msg.Strings{"foo", "bar"}
   157  			b, _ := a.MarshalMsg(nil)
   158  			if err := Send(rw, 1, b); err != nil {
   159  				t.Errorf("write error: %v", err)
   160  			}
   161  			return nil
   162  		},
   163  	}
   164  	closer, rw, _, _ := testPeer([]Protocol{proto})
   165  	defer closer()
   166  
   167  	if err := ExpectMsg(rw, 17, msg.Strings{"foo", "bar"}); err != nil {
   168  		t.Error(err)
   169  	}
   170  }
   171  
   172  func TestPeerPing(t *testing.T) {
   173  	closer, rw, _, _ := testPeer(nil)
   174  	defer closer()
   175  	if err := Send(rw, pingMsg, nil); err != nil {
   176  		t.Fatal(err)
   177  	}
   178  	if err := ExpectMsg(rw, pongMsg, nil); err != nil {
   179  		t.Error(err)
   180  	}
   181  }
   182  
   183  func TestPeerDisconnect(t *testing.T) {
   184  	logrus.Debug("started here")
   185  	closer, rw, _, disc := testPeer(nil)
   186  	defer closer()
   187  	b, _ := DiscQuitting.MarshalMsg(nil)
   188  	if err := Send(rw, discMsg, b); err != nil {
   189  		t.Fatal(err)
   190  	}
   191  	select {
   192  	case reason := <-disc:
   193  		if reason != DiscQuitting {
   194  			t.Errorf("run returned wrong reason: got %v, want %v", reason, DiscQuitting)
   195  			logrus.Debug("run returned wrong reason: got %v, want %v", reason, DiscQuitting)
   196  		}
   197  	case <-time.After(500 * time.Millisecond):
   198  		t.Error("peer did not return")
   199  	}
   200  	logrus.Debug("end here")
   201  }
   202  
   203  // This test is supposed to verify that Sender can reliably handle
   204  // multiple causes of disconnection occurring at the same time.
   205  func TestPeerDisconnectRace(t *testing.T) {
   206  	maybe := func() bool { return rand.Intn(1) == 1 }
   207  
   208  	for i := 0; i < 1000; i++ {
   209  		protoclose := make(chan error)
   210  		protodisc := make(chan DiscReason)
   211  		closer, rw, p, disc := testPeer([]Protocol{
   212  			{
   213  				Name:   "closereq",
   214  				Run:    func(p *Peer, rw MsgReadWriter) error { return <-protoclose },
   215  				Length: 1,
   216  			},
   217  			{
   218  				Name:   "disconnect",
   219  				Run:    func(p *Peer, rw MsgReadWriter) error { p.Disconnect(<-protodisc); return nil },
   220  				Length: 1,
   221  			},
   222  		})
   223  
   224  		// Simulate incoming messages.
   225  		go Send(rw, baseProtocolLength+1, nil)
   226  		go Send(rw, baseProtocolLength+2, nil)
   227  		// Close the network connection.
   228  		go closer()
   229  		// Make protocol "closereq" return.
   230  		protoclose <- errors.New("protocol closed")
   231  		// Make protocol "disconnect" call peer.Disconnect
   232  		protodisc <- DiscAlreadyConnected
   233  		// In some cases, simulate something else calling peer.Disconnect.
   234  		if maybe() {
   235  			go p.Disconnect(DiscInvalidIdentity)
   236  		}
   237  		// In some cases, simulate remote requesting a disconnect.
   238  		if maybe() {
   239  			b, _ := DiscQuitting.MarshalMsg(nil)
   240  			go Send(rw, discMsg, b)
   241  		}
   242  
   243  		select {
   244  		case <-disc:
   245  		case <-time.After(2 * time.Second):
   246  			// Sender.run should return quickly. If it doesn't the Sender
   247  			// goroutines are probably deadlocked. Call panic in order to
   248  			// show the stacks.
   249  			panic("Sender.run took to long to return.")
   250  		}
   251  	}
   252  }
   253  
   254  func TestNewPeer(t *testing.T) {
   255  	name := "nodename"
   256  	caps := []Cap{{"foo", 2}, {"bar", 3}}
   257  	id := randomID()
   258  	p := NewPeer(id, name, caps)
   259  	if p.ID() != id {
   260  		t.Errorf("ID mismatch: got %v, expected %v", p.ID(), id)
   261  	}
   262  	if p.Name() != name {
   263  		t.Errorf("Name mismatch: got %v, expected %v", p.Name(), name)
   264  	}
   265  	if !reflect.DeepEqual(p.Caps(), caps) {
   266  		t.Errorf("Caps mismatch: got %v, expected %v", p.Caps(), caps)
   267  	}
   268  
   269  	p.Disconnect(DiscAlreadyConnected) // Should not hang
   270  }
   271  
   272  func TestMatchProtocols(t *testing.T) {
   273  	tests := []struct {
   274  		Remote []Cap
   275  		Local  []Protocol
   276  		Match  map[string]protoRW
   277  	}{
   278  		{
   279  			// No remote capabilities
   280  			Local: []Protocol{{Name: "a"}},
   281  		},
   282  		{
   283  			// No local protocols
   284  			Remote: []Cap{{Name: "a"}},
   285  		},
   286  		{
   287  			// No mutual protocols
   288  			Remote: []Cap{{Name: "a"}},
   289  			Local:  []Protocol{{Name: "b"}},
   290  		},
   291  		{
   292  			// Some matches, some differences
   293  			Remote: []Cap{{Name: "local"}, {Name: "match1"}, {Name: "match2"}},
   294  			Local:  []Protocol{{Name: "match1"}, {Name: "match2"}, {Name: "remote"}},
   295  			Match:  map[string]protoRW{"match1": {Protocol: Protocol{Name: "match1"}}, "match2": {Protocol: Protocol{Name: "match2"}}},
   296  		},
   297  		{
   298  			// Various alphabetical ordering
   299  			Remote: []Cap{{Name: "aa"}, {Name: "ab"}, {Name: "bb"}, {Name: "ba"}},
   300  			Local:  []Protocol{{Name: "ba"}, {Name: "bb"}, {Name: "ab"}, {Name: "aa"}},
   301  			Match:  map[string]protoRW{"aa": {Protocol: Protocol{Name: "aa"}}, "ab": {Protocol: Protocol{Name: "ab"}}, "ba": {Protocol: Protocol{Name: "ba"}}, "bb": {Protocol: Protocol{Name: "bb"}}},
   302  		},
   303  		{
   304  			// No mutual versions
   305  			Remote: []Cap{{Version: 1}},
   306  			Local:  []Protocol{{Version: 2}},
   307  		},
   308  		{
   309  			// Multiple versions, single common
   310  			Remote: []Cap{{Version: 1}, {Version: 2}},
   311  			Local:  []Protocol{{Version: 2}, {Version: 3}},
   312  			Match:  map[string]protoRW{"": {Protocol: Protocol{Version: 2}}},
   313  		},
   314  		{
   315  			// Multiple versions, multiple common
   316  			Remote: []Cap{{Version: 1}, {Version: 2}, {Version: 3}, {Version: 4}},
   317  			Local:  []Protocol{{Version: 2}, {Version: 3}},
   318  			Match:  map[string]protoRW{"": {Protocol: Protocol{Version: 3}}},
   319  		},
   320  		{
   321  			// Various version orderings
   322  			Remote: []Cap{{Version: 4}, {Version: 1}, {Version: 3}, {Version: 2}},
   323  			Local:  []Protocol{{Version: 2}, {Version: 3}, {Version: 1}},
   324  			Match:  map[string]protoRW{"": {Protocol: Protocol{Version: 3}}},
   325  		},
   326  		{
   327  			// Versions overriding sub-protocol lengths
   328  			Remote: []Cap{{Version: 1}, {Version: 2}, {Version: 3}, {Name: "a"}},
   329  			Local:  []Protocol{{Version: 1, Length: 1}, {Version: 2, Length: 2}, {Version: 3, Length: 3}, {Name: "a"}},
   330  			Match:  map[string]protoRW{"": {Protocol: Protocol{Version: 3}}, "a": {Protocol: Protocol{Name: "a"}, offset: 3}},
   331  		},
   332  	}
   333  
   334  	for i, tt := range tests {
   335  		result := matchProtocols(tt.Local, tt.Remote, nil)
   336  		if len(result) != len(tt.Match) {
   337  			t.Errorf("test %d: negotiation mismatch: have %v, want %v", i, len(result), len(tt.Match))
   338  			continue
   339  		}
   340  		// Make sure all negotiated protocols are needed and correct
   341  		for name, proto := range result {
   342  			match, ok := tt.Match[name]
   343  			if !ok {
   344  				t.Errorf("test %d, proto '%s': negotiated but shouldn't have", i, name)
   345  				continue
   346  			}
   347  			if proto.Name != match.Name {
   348  				t.Errorf("test %d, proto '%s': name mismatch: have %v, want %v", i, name, proto.Name, match.Name)
   349  			}
   350  			if proto.Version != match.Version {
   351  				t.Errorf("test %d, proto '%s': version mismatch: have %v, want %v", i, name, proto.Version, match.Version)
   352  			}
   353  			if proto.offset-baseProtocolLength != match.offset {
   354  				t.Errorf("test %d, proto '%s': offset mismatch: have %v, want %v", i, name, proto.offset-baseProtocolLength, match.offset)
   355  			}
   356  		}
   357  		// Make sure no protocols missed negotiation
   358  		for name := range tt.Match {
   359  			if _, ok := result[name]; !ok {
   360  				t.Errorf("test %d, proto '%s': not negotiated, should have", i, name)
   361  				continue
   362  			}
   363  		}
   364  	}
   365  }
   366  
   367  func init() {
   368  	logrus.SetLevel(logrus.TraceLevel)
   369  	fname := "log334.log"
   370  	f, err := os.Create(fname)
   371  	if err != nil {
   372  		os.Remove(fname)
   373  		f, err = os.Create(fname)
   374  	}
   375  	if err != nil {
   376  		panic(err)
   377  	}
   378  	logrus.SetOutput(f)
   379  	log = logrus.StandardLogger()
   380  	//filenameHook := filename.NewHook()
   381  	//filenameHook.Field = "line"
   382  	//logrus.AddHook(filenameHook)
   383  }