fixed merge issue

4 files changed, 263 insertions, 30 deletions

diff --git a/p2p/handshake.go b/p2p/handshake.go
index 79395f23f..8e611cfd5 100644
--- a/p2p/handshake.go
+++ b/p2p/handshake.go
@@ -70,21 +70,21 @@ type protoHandshake struct {
 // If dial is non-nil, the connection the local node is the initiator.
 // If atcap is true, the connection will be disconnected with DiscTooManyPeers
 // after the key exchange.
-func setupConn(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, dial *discover.Node, atcap bool) (*conn, error) {
+func setupConn(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, dial *discover.Node, atcap bool, trusted map[discover.NodeID]bool) (*conn, error) {
 	if dial == nil {
-		return setupInboundConn(fd, prv, our, atcap)
+		return setupInboundConn(fd, prv, our, atcap, trusted)
 	} else {
-		return setupOutboundConn(fd, prv, our, dial, atcap)
+		return setupOutboundConn(fd, prv, our, dial, atcap, trusted)
 	}
 }
 
-func setupInboundConn(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, atcap bool) (*conn, error) {
+func setupInboundConn(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, atcap bool, trusted map[discover.NodeID]bool) (*conn, error) {
 	secrets, err := receiverEncHandshake(fd, prv, nil)
 	if err != nil {
 		return nil, fmt.Errorf("encryption handshake failed: %v", err)
 	}
 	rw := newRlpxFrameRW(fd, secrets)
-	if atcap {
+	if atcap && !trusted[secrets.RemoteID] {
 		SendItems(rw, discMsg, DiscTooManyPeers)
 		return nil, errors.New("we have too many peers")
 	}
@@ -99,13 +99,13 @@ func setupInboundConn(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, a
 	return &conn{rw, rhs}, nil
 }
 
-func setupOutboundConn(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, dial *discover.Node, atcap bool) (*conn, error) {
+func setupOutboundConn(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, dial *discover.Node, atcap bool, trusted map[discover.NodeID]bool) (*conn, error) {
 	secrets, err := initiatorEncHandshake(fd, prv, dial.ID, nil)
 	if err != nil {
 		return nil, fmt.Errorf("encryption handshake failed: %v", err)
 	}
 	rw := newRlpxFrameRW(fd, secrets)
-	if atcap {
+	if atcap && !trusted[secrets.RemoteID] {
 		SendItems(rw, discMsg, DiscTooManyPeers)
 		return nil, errors.New("we have too many peers")
 	}
diff --git a/p2p/handshake_test.go b/p2p/handshake_test.go
index f618ef20d..9018e14f2 100644
--- a/p2p/handshake_test.go
+++ b/p2p/handshake_test.go
@@ -143,7 +143,7 @@ func TestSetupConn(t *testing.T) {
 	done := make(chan struct{})
 	go func() {
 		defer close(done)
-		conn0, err := setupConn(fd0, prv0, hs0, node1, false)
+		conn0, err := setupConn(fd0, prv0, hs0, node1, false, nil)
 		if err != nil {
 			t.Errorf("outbound side error: %v", err)
 			return
@@ -156,7 +156,7 @@ func TestSetupConn(t *testing.T) {
 		}
 	}()
 
-	conn1, err := setupConn(fd1, prv1, hs1, nil, false)
+	conn1, err := setupConn(fd1, prv1, hs1, nil, false, nil)
 	if err != nil {
 		t.Fatalf("inbound side error: %v", err)
 	}
diff --git a/p2p/server.go b/p2p/server.go
index e648c72c9..4a78e0bb9 100644
--- a/p2p/server.go
+++ b/p2p/server.go
@@ -18,8 +18,9 @@ import (
 )
 
 const (
-	defaultDialTimeout   = 10 * time.Second
-	refreshPeersInterval = 30 * time.Second
+	defaultDialTimeout      = 10 * time.Second
+	refreshPeersInterval    = 30 * time.Second
+	staticPeerCheckInterval = 15 * time.Second
 
 	// This is the maximum number of inbound connection
 	// that are allowed to linger between 'accepted' and
@@ -59,6 +60,14 @@ type Server struct {
 	// with the rest of the network.
 	BootstrapNodes []*discover.Node
 
+	// Static nodes are used as pre-configured connections which are always
+	// maintained and re-connected on disconnects.
+	StaticNodes []*discover.Node
+
+	// Trusted nodes are used as pre-configured connections which are always
+	// allowed to connect, even above the peer limit.
+	TrustedNodes []*discover.Node
+
 	// NodeDatabase is the path to the database containing the previously seen
 	// live nodes in the network.
 	NodeDatabase string
@@ -95,20 +104,23 @@ type Server struct {
 
 	ourHandshake *protoHandshake
 
-	lock    sync.RWMutex // protects running and peers
-	running bool
-	peers   map[discover.NodeID]*Peer
+	lock         sync.RWMutex // protects running, peers and the trust fields
+	running      bool
+	peers        map[discover.NodeID]*Peer
+	staticNodes  map[discover.NodeID]*discover.Node // Map of currently maintained static remote nodes
+	staticDial   chan *discover.Node                // Dial request channel reserved for the static nodes
+	staticCycle  time.Duration                      // Overrides staticPeerCheckInterval, used for testing
+	trustedNodes map[discover.NodeID]bool           // Set of currently trusted remote nodes
 
 	ntab     *discover.Table
 	listener net.Listener
 
-	quit        chan struct{}
-	loopWG      sync.WaitGroup // {dial,listen,nat}Loop
-	peerWG      sync.WaitGroup // active peer goroutines
-	peerConnect chan *discover.Node
+	quit   chan struct{}
+	loopWG sync.WaitGroup // {dial,listen,nat}Loop
+	peerWG sync.WaitGroup // active peer goroutines
 }
 
-type setupFunc func(net.Conn, *ecdsa.PrivateKey, *protoHandshake, *discover.Node, bool) (*conn, error)
+type setupFunc func(net.Conn, *ecdsa.PrivateKey, *protoHandshake, *discover.Node, bool, map[discover.NodeID]bool) (*conn, error)
 type newPeerHook func(*Peer)
 
 // Peers returns all connected peers.
@@ -131,10 +143,14 @@ func (srv *Server) PeerCount() int {
 	return n
 }
 
-// SuggestPeer creates a connection to the given Node if it
-// is not already connected.
-func (srv *Server) SuggestPeer(n *discover.Node) {
-	srv.peerConnect <- n
+// AddPeer connects to the given node and maintains the connection until the
+// server is shut down. If the connection fails for any reason, the server will
+// attempt to reconnect the peer.
+func (srv *Server) AddPeer(node *discover.Node) {
+	srv.lock.Lock()
+	defer srv.lock.Unlock()
+
+	srv.staticNodes[node.ID] = node
 }
 
 // Broadcast sends an RLP-encoded message to all connected peers.
@@ -195,7 +211,18 @@ func (srv *Server) Start() (err error) {
 	}
 	srv.quit = make(chan struct{})
 	srv.peers = make(map[discover.NodeID]*Peer)
-	srv.peerConnect = make(chan *discover.Node)
+
+	// Create the current trust maps, and the associated dialing channel
+	srv.trustedNodes = make(map[discover.NodeID]bool)
+	for _, node := range srv.TrustedNodes {
+		srv.trustedNodes[node.ID] = true
+	}
+	srv.staticNodes = make(map[discover.NodeID]*discover.Node)
+	for _, node := range srv.StaticNodes {
+		srv.staticNodes[node.ID] = node
+	}
+	srv.staticDial = make(chan *discover.Node)
+
 	if srv.setupFunc == nil {
 		srv.setupFunc = setupConn
 	}
@@ -229,6 +256,8 @@ func (srv *Server) Start() (err error) {
 	if srv.NoDial && srv.ListenAddr == "" {
 		glog.V(logger.Warn).Infoln("I will be kind-of useless, neither dialing nor listening.")
 	}
+	// maintain the static peers
+	go srv.staticNodesLoop()
 
 	srv.running = true
 	return nil
@@ -323,6 +352,45 @@ func (srv *Server) listenLoop() {
 	}
 }
 
+// staticNodesLoop is responsible for periodically checking that static
+// connections are actually live, and requests dialing if not.
+func (srv *Server) staticNodesLoop() {
+	// Create a default maintenance ticker, but override it requested
+	cycle := staticPeerCheckInterval
+	if srv.staticCycle != 0 {
+		cycle = srv.staticCycle
+	}
+	tick := time.NewTicker(cycle)
+
+	for {
+		select {
+		case <-srv.quit:
+			return
+
+		case <-tick.C:
+			// Collect all the non-connected static nodes
+			needed := []*discover.Node{}
+			srv.lock.RLock()
+			for id, node := range srv.staticNodes {
+				if _, ok := srv.peers[id]; !ok {
+					needed = append(needed, node)
+				}
+			}
+			srv.lock.RUnlock()
+
+			// Try to dial each of them (don't hang if server terminates)
+			for _, node := range needed {
+				glog.V(logger.Debug).Infof("Dialing static peer %v", node)
+				select {
+				case srv.staticDial <- node:
+				case <-srv.quit:
+					return
+				}
+			}
+		}
+	}
+}
+
 func (srv *Server) dialLoop() {
 	var (
 		dialed      = make(chan *discover.Node)
@@ -373,7 +441,7 @@ func (srv *Server) dialLoop() {
 				// below MaxPeers.
 				refresh.Reset(refreshPeersInterval)
 			}
-		case dest := <-srv.peerConnect:
+		case dest := <-srv.staticDial:
 			dial(dest)
 		case dests := <-findresults:
 			for _, dest := range dests {
@@ -416,10 +484,18 @@ func (srv *Server) startPeer(fd net.Conn, dest *discover.Node) {
 	// returns during that exchange need to call peerWG.Done because
 	// the callers of startPeer added the peer to the wait group already.
 	fd.SetDeadline(time.Now().Add(handshakeTimeout))
+
+	// Check capacity, but override for static nodes
 	srv.lock.RLock()
 	atcap := len(srv.peers) == srv.MaxPeers
+	if dest != nil {
+		if _, ok := srv.staticNodes[dest.ID]; ok {
+			atcap = false
+		}
+	}
 	srv.lock.RUnlock()
-	conn, err := srv.setupFunc(fd, srv.PrivateKey, srv.ourHandshake, dest, atcap)
+
+	conn, err := srv.setupFunc(fd, srv.PrivateKey, srv.ourHandshake, dest, atcap, srv.trustedNodes)
 	if err != nil {
 		fd.Close()
 		glog.V(logger.Debug).Infof("Handshake with %v failed: %v", fd.RemoteAddr(), err)
@@ -472,11 +548,18 @@ func (srv *Server) addPeer(id discover.NodeID, p *Peer) (bool, DiscReason) {
 	return true, 0
 }
 
+// checkPeer verifies whether a peer looks promising and should be allowed/kept
+// in the pool, or if it's of no use.
 func (srv *Server) checkPeer(id discover.NodeID) (bool, DiscReason) {
+	// First up, figure out if the peer is static or trusted
+	_, static := srv.staticNodes[id]
+	trusted := srv.trustedNodes[id]
+
+	// Make sure the peer passes all required checks
 	switch {
 	case !srv.running:
 		return false, DiscQuitting
-	case len(srv.peers) >= srv.MaxPeers:
+	case !static && !trusted && len(srv.peers) >= srv.MaxPeers:
 		return false, DiscTooManyPeers
 	case srv.peers[id] != nil:
 		return false, DiscAlreadyConnected
diff --git a/p2p/server_test.go b/p2p/server_test.go
index 86514b650..3f9db343c 100644
--- a/p2p/server_test.go
+++ b/p2p/server_test.go
@@ -22,7 +22,7 @@ func startTestServer(t *testing.T, pf newPeerHook) *Server {
 		ListenAddr:  "127.0.0.1:0",
 		PrivateKey:  newkey(),
 		newPeerHook: pf,
-		setupFunc: func(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, dial *discover.Node, atcap bool) (*conn, error) {
+		setupFunc: func(fd net.Conn, prv *ecdsa.PrivateKey, our *protoHandshake, dial *discover.Node, atcap bool, trusted map[discover.NodeID]bool) (*conn, error) {
 			id := randomID()
 			rw := newRlpxFrameRW(fd, secrets{
 				MAC:        zero16,
@@ -102,7 +102,7 @@ func TestServerDial(t *testing.T) {
 
 	// tell the server to connect
 	tcpAddr := listener.Addr().(*net.TCPAddr)
-	srv.SuggestPeer(&discover.Node{IP: tcpAddr.IP, TCP: uint16(tcpAddr.Port)})
+	srv.staticDial <- &discover.Node{IP: tcpAddr.IP, TCP: uint16(tcpAddr.Port)}
 
 	select {
 	case conn := <-accepted:
@@ -200,7 +200,7 @@ func TestServerDisconnectAtCap(t *testing.T) {
 		// 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(), false)
+		_, err = setupConn(conn, key, hs, srv.Self(), false, srv.trustedNodes)
 		if i == nconns-1 {
 			// When handling the last connection, the server should
 			// disconnect immediately instead of running the protocol
@@ -219,6 +219,156 @@ func TestServerDisconnectAtCap(t *testing.T) {
 	}
 }
 
+// Tests that static peers are (re)connected, and done so even above max peers.
+func TestServerStaticPeers(t *testing.T) {
+	defer testlog(t).detach()
+
+	// 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(), false, server.trustedNodes); 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) {
+	defer testlog(t).detach()
+
+	// 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(), false, server.trustedNodes); 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(), false, server.trustedNodes); 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")
+	}
+}
+
 func newkey() *ecdsa.PrivateKey {
 	key, err := crypto.GenerateKey()
 	if err != nil {