parent
525b20ae20
commit
0e4568253a
@ -0,0 +1,33 @@ |
||||
package blockedpeers |
||||
|
||||
import ( |
||||
"github.com/harmony-one/harmony/internal/utils/lrucache" |
||||
libp2p_peer "github.com/libp2p/go-libp2p/core/peer" |
||||
"time" |
||||
) |
||||
|
||||
type Manager struct { |
||||
internal *lrucache.Cache[libp2p_peer.ID, time.Time] |
||||
} |
||||
|
||||
func NewManager(size int) *Manager { |
||||
return &Manager{ |
||||
internal: lrucache.NewCache[libp2p_peer.ID, time.Time](size), |
||||
} |
||||
} |
||||
|
||||
func (m *Manager) IsBanned(key libp2p_peer.ID, now time.Time) bool { |
||||
future, ok := m.internal.Get(key) |
||||
if ok { |
||||
return future.After(now) // future > now
|
||||
} |
||||
return ok |
||||
} |
||||
|
||||
func (m *Manager) Ban(key libp2p_peer.ID, future time.Time) { |
||||
m.internal.Set(key, future) |
||||
} |
||||
|
||||
func (m *Manager) Contains(key libp2p_peer.ID) bool { |
||||
return m.internal.Contains(key) |
||||
} |
@ -0,0 +1,26 @@ |
||||
package blockedpeers |
||||
|
||||
import ( |
||||
libp2p_peer "github.com/libp2p/go-libp2p/core/peer" |
||||
"github.com/stretchr/testify/require" |
||||
"testing" |
||||
"time" |
||||
) |
||||
|
||||
func TestNewManager(t *testing.T) { |
||||
var ( |
||||
peer1 libp2p_peer.ID = "peer1" |
||||
now = time.Now() |
||||
m = NewManager(4) |
||||
) |
||||
|
||||
t.Run("check_empty", func(t *testing.T) { |
||||
require.False(t, m.IsBanned(peer1, now), "peer1 should not be banned") |
||||
}) |
||||
t.Run("ban_peer1", func(t *testing.T) { |
||||
m.Ban(peer1, now.Add(2*time.Second)) |
||||
require.True(t, m.IsBanned(peer1, now), "peer1 should be banned") |
||||
require.False(t, m.IsBanned(peer1, now.Add(3*time.Second)), "peer1 should not be banned after 3 seconds") |
||||
}) |
||||
|
||||
} |
Loading…
Reference in new issue