diff options
-rw-r--r-- | p2p/peer.go | 12 | ||||
-rw-r--r-- | p2p/server.go | 13 |
2 files changed, 21 insertions, 4 deletions
diff --git a/p2p/peer.go b/p2p/peer.go index cdf9ba965..ac691f2ce 100644 --- a/p2p/peer.go +++ b/p2p/peer.go @@ -211,6 +211,18 @@ func (p *Peer) handle(msg Msg) error { return nil } +func countMatchingProtocols(protocols []Protocol, caps []Cap) int { + n := 0 + for _, cap := range caps { + for _, proto := range protocols { + if proto.Name == cap.Name && proto.Version == cap.Version { + n++ + } + } + } + return n +} + // matchProtocols creates structures for matching named subprotocols. func matchProtocols(protocols []Protocol, caps []Cap, rw MsgReadWriter) map[string]*protoRW { sort.Sort(capsByName(caps)) diff --git a/p2p/server.go b/p2p/server.go index 171798a1d..3c6fb5893 100644 --- a/p2p/server.go +++ b/p2p/server.go @@ -518,7 +518,7 @@ func (srv *Server) startPeer(fd net.Conn, dest *discover.Node) { conn: fd, rtimeout: frameReadTimeout, wtimeout: frameWriteTimeout, } p := newPeer(fd, conn, srv.Protocols) - if ok, reason := srv.addPeer(conn.ID, p); !ok { + if ok, reason := srv.addPeer(conn, p); !ok { glog.V(logger.Detail).Infof("Not adding %v (%v)\n", p, reason) p.politeDisconnect(reason) srv.peerWG.Done() @@ -564,13 +564,18 @@ func (srv *Server) runPeer(p *Peer) { }) } -func (srv *Server) addPeer(id discover.NodeID, p *Peer) (bool, DiscReason) { +func (srv *Server) addPeer(conn *conn, p *Peer) (bool, DiscReason) { + // drop connections with no matching protocols. + if len(srv.Protocols) > 0 && countMatchingProtocols(srv.Protocols, conn.protoHandshake.Caps) == 0 { + return false, DiscUselessPeer + } + // add the peer if it passes the other checks. srv.lock.Lock() defer srv.lock.Unlock() - if ok, reason := srv.checkPeer(id); !ok { + if ok, reason := srv.checkPeer(conn.ID); !ok { return false, reason } - srv.peers[id] = p + srv.peers[conn.ID] = p return true, 0 } |