forked from zjshen14/go-p2p
-
Notifications
You must be signed in to change notification settings - Fork 5
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
5 changed files
with
195 additions
and
41 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,72 @@ | ||
| package p2p | ||
|
|
||
| import ( | ||
| "time" | ||
|
|
||
| "github.com/iotexproject/go-pkgs/cache" | ||
| "github.com/libp2p/go-libp2p-core" | ||
| ) | ||
|
|
||
| // CountBlocklist is a count-based blacklist implementation. Entry is being added | ||
| // to an LRU cache, once the count reaches a threshold the entry is blocked for a | ||
| // certain timeout. The entry is then unblocked once the timeout expires | ||
| type CountBlocklist struct { | ||
| count int // number of times an entry is added before being blocked | ||
| ttl time.Duration // timeout an entry will be blocked | ||
| counter *cache.ThreadSafeLruCache | ||
| timeout *cache.ThreadSafeLruCache | ||
| } | ||
|
|
||
| // NewCountBlocklist creates a new CountBlocklist | ||
| func NewCountBlocklist(cap, count int, ttl time.Duration) *CountBlocklist { | ||
| return &CountBlocklist{ | ||
| count: count, | ||
| ttl: ttl, | ||
| counter: cache.NewThreadSafeLruCache(cap), | ||
| timeout: cache.NewThreadSafeLruCache(cap), | ||
| } | ||
| } | ||
|
|
||
| // Blocked returns true if the name is blocked | ||
| func (b *CountBlocklist) Blocked(name core.PeerID, t time.Time) bool { | ||
| v, ok := b.timeout.Get(name) | ||
| if !ok { | ||
| return false | ||
| } | ||
|
|
||
| if v.(time.Time).After(t) { | ||
| return true | ||
| } | ||
|
|
||
| // timeout passed, remove name off the blocklist | ||
| b.remove(name) | ||
| return false | ||
| } | ||
|
|
||
| // Add tries to add the name to blocklist | ||
| func (b *CountBlocklist) Add(name core.PeerID, t time.Time) { | ||
| var count int | ||
| if v, ok := b.counter.Get(name); !ok { | ||
| count = 1 | ||
| } else { | ||
| count = v.(int) + 1 | ||
| } | ||
| b.counter.Add(name, count) | ||
|
|
||
| // add to blocklist once reaching count | ||
| if count >= b.count { | ||
| b.timeout.Add(name, t.Add(b.ttl)) | ||
| } | ||
| } | ||
|
|
||
| // Clear clears the blocklist | ||
| func (b *CountBlocklist) Clear() { | ||
| b.timeout.Clear() | ||
| b.counter.Clear() | ||
| } | ||
|
|
||
| // remove takes name off the blocklist | ||
| func (b *CountBlocklist) remove(name core.PeerID) { | ||
| b.counter.Remove(name) | ||
| b.timeout.Remove(name) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,48 @@ | ||
| package p2p | ||
|
|
||
| import ( | ||
| "testing" | ||
| "time" | ||
|
|
||
| "github.com/libp2p/go-libp2p-core" | ||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| func TestBlockList(t *testing.T) { | ||
| r := require.New(t) | ||
|
|
||
| cfg := DefaultConfig | ||
| now := time.Now() | ||
| name := core.PeerID("alfa") | ||
| withinBlockTTL := now.Add(cfg.BlockListCleanupInterval / 2) | ||
| pastBlockTTL := now.Add(cfg.BlockListCleanupInterval * 2) | ||
|
|
||
| blockTests := []struct { | ||
| curTime time.Time | ||
| blocked bool | ||
| }{ | ||
| {now, false}, | ||
| {now, false}, | ||
| {withinBlockTTL, true}, | ||
| {withinBlockTTL, true}, | ||
| {pastBlockTTL, false}, | ||
| {withinBlockTTL, false}, | ||
| {withinBlockTTL, false}, | ||
| {pastBlockTTL, false}, | ||
| {now, false}, | ||
| {now, false}, | ||
| {withinBlockTTL, true}, | ||
| } | ||
|
|
||
| list := NewCountBlocklist(cfg.BlockListLRUSize, cfg.BlockListErrorThreshold, cfg.BlockListCleanupInterval) | ||
| for _, v := range blockTests { | ||
| list.Add(name, now) | ||
| r.Equal(v.blocked, list.Blocked(name, v.curTime)) | ||
| } | ||
| r.Equal(1, list.counter.Len()) | ||
| r.Equal(1, list.timeout.Len()) | ||
|
|
||
| list.Clear() | ||
| r.Zero(list.counter.Len()) | ||
| r.Zero(list.timeout.Len()) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters