fix ban node failed (#256)

[bytom/vapor.git] / netsync / peers / peer.go

diff --git a/netsync/peers/peer.go b/netsync/peers/peer.go
index d0eb7ee..e74a458 100644 (file)
--- a/netsync/peers/peer.go
+++ b/netsync/peers/peer.go
@@ -36,6 +36,7 @@ var (
 type BasePeer interface {
        Addr() net.Addr
        ID() string
+       RemoteAddrHost() string
        ServiceFlag() consensus.ServiceFlag
        TrafficStatus() (*flowrate.Status, *flowrate.Status)
        TrySend(byte, interface{}) bool
@@ -45,7 +46,7 @@ type BasePeer interface {
 //BasePeerSet is the intergace for connection level peer manager
 type BasePeerSet interface {
        StopPeerGracefully(string)
-       IsBanned(peerID string, level byte, reason string) bool
+       IsBanned(ip string, level byte, reason string) bool
 }
 
 type BroadcastMsg interface {
@@ -417,7 +418,8 @@ func (ps *PeerSet) ProcessIllegal(peerID string, level byte, reason string) {
        if peer == nil {
                return
        }
-       if banned := ps.IsBanned(peer.Addr().String(), level, reason); banned {
+
+       if banned := ps.IsBanned(peer.RemoteAddrHost(), level, reason); banned {
                ps.RemovePeer(peerID)
        }
        return