From b08d945d7c62b6dc861d02fe490dc1cf5767ae6e Mon Sep 17 00:00:00 2001 From: 9seconds Date: Mon, 15 Mar 2021 21:31:07 +0300 Subject: [PATCH] Add event stream module --- antireplay/noop.go | 1 + antireplay/noop_test.go | 2 + antireplay/stable_bloom_filter.go | 2 + antireplay/stable_bloom_filter_test.go | 2 + events/event_stream.go | 64 ++++++++++++++++++++++++++ events/init.go | 12 +++++ events/multi_observer.go | 59 ++++++++++++++++++++++++ ipblocklist/firehol.go | 33 +++++++++---- ipblocklist/firehol_test.go | 21 +++++++-- ipblocklist/noop.go | 1 + mtglib/events.go | 25 ++++++++++ mtglib/init.go | 11 +++++ testlib/events_observer_mock.go | 7 +++ 13 files changed, 225 insertions(+), 15 deletions(-) create mode 100644 events/event_stream.go create mode 100644 events/init.go create mode 100644 events/multi_observer.go create mode 100644 mtglib/events.go create mode 100644 testlib/events_observer_mock.go diff --git a/antireplay/noop.go b/antireplay/noop.go index 8613e40..b07d34f 100644 --- a/antireplay/noop.go +++ b/antireplay/noop.go @@ -5,6 +5,7 @@ import "github.com/9seconds/mtg/v2/mtglib" type noop struct{} func (n noop) SeenBefore(_ []byte) bool { return false } +func (n noop) Shutdown() {} func NewNoop() mtglib.AntiReplayCache { return noop{} diff --git a/antireplay/noop_test.go b/antireplay/noop_test.go index b78c48d..42f39eb 100644 --- a/antireplay/noop_test.go +++ b/antireplay/noop_test.go @@ -18,6 +18,8 @@ func (suite *NoopTestSuite) TestOp() { suite.False(filter.SeenBefore([]byte{4, 5, 6})) suite.False(filter.SeenBefore([]byte{1, 2, 3})) suite.False(filter.SeenBefore([]byte{4, 5, 6})) + + filter.Shutdown() } func TestNoop(t *testing.T) { diff --git a/antireplay/stable_bloom_filter.go b/antireplay/stable_bloom_filter.go index b52631e..ea4f5a7 100644 --- a/antireplay/stable_bloom_filter.go +++ b/antireplay/stable_bloom_filter.go @@ -20,6 +20,8 @@ func (s *stableBloomFilter) SeenBefore(digest []byte) bool { return s.filter.TestAndAdd(digest) } +func (s *stableBloomFilter) Shutdown() {} + func NewStableBloomFilter(byteSize uint, errorRate float64) mtglib.AntiReplayCache { sf := boom.NewDefaultStableBloomFilter(byteSize*8, errorRate) // nolint: gomnd sf.SetHash(xxhash.New64()) diff --git a/antireplay/stable_bloom_filter_test.go b/antireplay/stable_bloom_filter_test.go index accb590..9b2517e 100644 --- a/antireplay/stable_bloom_filter_test.go +++ b/antireplay/stable_bloom_filter_test.go @@ -18,6 +18,8 @@ func (suite *StableBloomFilterTestSuite) TestOp() { suite.False(filter.SeenBefore([]byte{4, 5, 6})) suite.True(filter.SeenBefore([]byte{1, 2, 3})) suite.True(filter.SeenBefore([]byte{4, 5, 6})) + + filter.Shutdown() } func TestStableBloomFilter(t *testing.T) { diff --git a/events/event_stream.go b/events/event_stream.go new file mode 100644 index 0000000..577e713 --- /dev/null +++ b/events/event_stream.go @@ -0,0 +1,64 @@ +package events + +import ( + "context" + "runtime" + + "github.com/9seconds/mtg/v2/mtglib" + "github.com/OneOfOne/xxhash" +) + +type eventStream struct { + ctx context.Context + ctxCancel context.CancelFunc + chans []chan mtglib.Event +} + +func (e eventStream) Send(ctx context.Context, evt mtglib.Event) { + chanNo := int(xxhash.ChecksumString32(evt.ConnectionID())) % len(e.chans) + + select { + case <-ctx.Done(): + case <-e.ctx.Done(): + case e.chans[chanNo] <- evt: + } +} + +func (e eventStream) Shutdown() { + e.ctxCancel() +} + +func NewEventStream(observerFactories []ObserverFactory) mtglib.EventStream { + ctx, cancel := context.WithCancel(context.Background()) + rv := eventStream{ + ctx: ctx, + ctxCancel: cancel, + chans: make([]chan mtglib.Event, runtime.NumCPU()), + } + + for i := 0; i < runtime.NumCPU(); i++ { + rv.chans[i] = make(chan mtglib.Event, 1) + + go eventStreamProcessor(ctx, rv.chans[i], newMultiObserver(observerFactories)) + } + + return rv +} + +func eventStreamProcessor(ctx context.Context, eventChan <-chan mtglib.Event, observer Observer) { + defer observer.Shutdown() + + for { + select { + case <-ctx.Done(): + return + case evt := <-eventChan: + switch typedEvt := evt.(type) { + case mtglib.EventStart: + observer.EventStart(typedEvt) + case mtglib.EventFinish: + observer.EventFinish(typedEvt) + } + } + } +} diff --git a/events/init.go b/events/init.go new file mode 100644 index 0000000..f84b03b --- /dev/null +++ b/events/init.go @@ -0,0 +1,12 @@ +package events + +import "github.com/9seconds/mtg/v2/mtglib" + +type Observer interface { + EventStart(mtglib.EventStart) + EventFinish(mtglib.EventFinish) + + Shutdown() +} + +type ObserverFactory func() Observer diff --git a/events/multi_observer.go b/events/multi_observer.go new file mode 100644 index 0000000..9370fb0 --- /dev/null +++ b/events/multi_observer.go @@ -0,0 +1,59 @@ +package events + +import ( + "sync" + + "github.com/9seconds/mtg/v2/mtglib" +) + +type multiObserver struct { + observers []Observer +} + +func (m multiObserver) EventStart(evt mtglib.EventStart) { + wg := &sync.WaitGroup{} + wg.Add(len(m.observers)) + + for _, v := range m.observers { + go func(obs Observer) { + defer wg.Done() + + obs.EventStart(evt) + }(v) + } + + wg.Wait() +} + +func (m multiObserver) EventFinish(evt mtglib.EventFinish) { + wg := &sync.WaitGroup{} + wg.Add(len(m.observers)) + + for _, v := range m.observers { + go func(obs Observer) { + defer wg.Done() + + obs.EventFinish(evt) + }(v) + } + + wg.Wait() +} + +func (m multiObserver) Shutdown() { + for _, v := range m.observers { + v.Shutdown() + } +} + +func newMultiObserver(factories []ObserverFactory) Observer { + observers := make([]Observer, len(factories)) + + for i, v := range factories { + observers[i] = v() + } + + return multiObserver{ + observers: observers, + } +} diff --git a/ipblocklist/firehol.go b/ipblocklist/firehol.go index e4c0370..076ab6e 100644 --- a/ipblocklist/firehol.go +++ b/ipblocklist/firehol.go @@ -29,14 +29,20 @@ const ( var fireholRegexpComment = regexp.MustCompile(`\s*#.*?$`) type Firehol struct { - logger mtglib.Logger - rwMutex sync.RWMutex + ctx context.Context + ctxCancel context.CancelFunc + logger mtglib.Logger + + rwMutex sync.RWMutex + remoteURLs []string localFiles []string + httpClient *http.Client workerPool *ants.Pool - treeV4 *bool_tree.TreeV4 - treeV6 *bool_tree.TreeV6 + + treeV4 *bool_tree.TreeV4 + treeV6 *bool_tree.TreeV6 } func (f *Firehol) Contains(ip net.IP) bool { @@ -76,7 +82,7 @@ func (f *Firehol) containsIPv6(addr net.IP) bool { return false } -func (f *Firehol) Run(ctx context.Context, updateEach time.Duration) { +func (f *Firehol) Run(updateEach time.Duration) { ticker := time.NewTicker(updateEach) defer func() { @@ -88,24 +94,28 @@ func (f *Firehol) Run(ctx context.Context, updateEach time.Duration) { } }() - if err := f.update(ctx); err != nil { + if err := f.update(); err != nil { f.logger.WarningError("cannot update blocklist", err) } for { select { - case <-ctx.Done(): + case <-f.ctx.Done(): return case <-ticker.C: - if err := f.update(ctx); err != nil { + if err := f.update(); err != nil { f.logger.WarningError("cannot update blocklist", err) } } } } -func (f *Firehol) update(ctx context.Context) error { // nolint: funlen, cyclop - ctx, cancel := context.WithCancel(ctx) +func (f *Firehol) Shutdown() { + f.ctxCancel() +} + +func (f *Firehol) update() error { // nolint: funlen, cyclop + ctx, cancel := context.WithCancel(f.ctx) defer cancel() wg := &sync.WaitGroup{} @@ -314,8 +324,11 @@ func NewFirehol(logger mtglib.Logger, network mtglib.Network, } workerPool, _ := ants.NewPool(int(downloadConcurrency)) + ctx, cancel := context.WithCancel(context.Background()) return &Firehol{ + ctx: ctx, + ctxCancel: cancel, logger: logger.Named("firehol"), httpClient: network.MakeHTTPClient(nil), treeV4: bool_tree.NewTreeV4(), diff --git a/ipblocklist/firehol_test.go b/ipblocklist/firehol_test.go index 5b18820..ad353fe 100644 --- a/ipblocklist/firehol_test.go +++ b/ipblocklist/firehol_test.go @@ -1,7 +1,6 @@ package ipblocklist_test import ( - "context" "io" "net" "net/http" @@ -72,12 +71,15 @@ func (suite *FireholTestSuite) TestLocalFail() { suite.NoError(err) - go blocklist.Run(context.Background(), time.Hour) + go blocklist.Run(time.Hour) time.Sleep(500 * time.Millisecond) suite.False(blocklist.Contains(net.ParseIP("10.0.0.10"))) suite.False(blocklist.Contains(net.ParseIP("127.0.0.1"))) + + blocklist.Shutdown() + time.Sleep(500 * time.Millisecond) } func (suite *FireholTestSuite) TestLocalOk() { @@ -87,12 +89,15 @@ func (suite *FireholTestSuite) TestLocalOk() { suite.NoError(err) - go blocklist.Run(context.Background(), time.Hour) + go blocklist.Run(time.Hour) time.Sleep(500 * time.Millisecond) suite.True(blocklist.Contains(net.ParseIP("10.0.0.10"))) suite.False(blocklist.Contains(net.ParseIP("127.0.0.1"))) + + blocklist.Shutdown() + time.Sleep(500 * time.Millisecond) } func (suite *FireholTestSuite) TestRemoteFail() { @@ -102,11 +107,14 @@ func (suite *FireholTestSuite) TestRemoteFail() { suite.NoError(err) - go blocklist.Run(context.Background(), time.Hour) + go blocklist.Run(time.Hour) time.Sleep(500 * time.Millisecond) suite.False(blocklist.Contains(net.ParseIP("10.2.2.2"))) + + blocklist.Shutdown() + time.Sleep(500 * time.Millisecond) } func (suite *FireholTestSuite) TestMixed() { @@ -123,12 +131,15 @@ func (suite *FireholTestSuite) TestMixed() { suite.NoError(err) - go blocklist.Run(context.Background(), time.Hour) + go blocklist.Run(time.Hour) time.Sleep(500 * time.Millisecond) suite.True(blocklist.Contains(net.ParseIP("10.2.2.2"))) suite.True(blocklist.Contains(net.ParseIP("10.1.0.100"))) + + blocklist.Shutdown() + time.Sleep(500 * time.Millisecond) } func TestFirehol(t *testing.T) { diff --git a/ipblocklist/noop.go b/ipblocklist/noop.go index dcbed19..4b29449 100644 --- a/ipblocklist/noop.go +++ b/ipblocklist/noop.go @@ -9,6 +9,7 @@ import ( type noop struct{} func (n noop) Contains(ip net.IP) bool { return false } +func (n noop) Shutdown() {} func NewNoop() mtglib.IPBlocklist { return noop{} diff --git a/mtglib/events.go b/mtglib/events.go new file mode 100644 index 0000000..ef361a2 --- /dev/null +++ b/mtglib/events.go @@ -0,0 +1,25 @@ +package mtglib + +import ( + "net" + "time" +) + +type eventBase struct { + CreatedAt time.Time + ConnID string +} + +func (e eventBase) ConnectionID() string { + return e.ConnID +} + +type EventStart struct { + eventBase + + RemoteIP net.IP +} + +type EventFinish struct { + eventBase +} diff --git a/mtglib/init.go b/mtglib/init.go index 12e1393..8a3b764 100644 --- a/mtglib/init.go +++ b/mtglib/init.go @@ -19,10 +19,21 @@ type Network interface { type AntiReplayCache interface { SeenBefore(data []byte) bool + Shutdown() } type IPBlocklist interface { Contains(net.IP) bool + Shutdown() +} + +type Event interface { + ConnectionID() string +} + +type EventStream interface { + Send(context.Context, Event) + Shutdown() } type Logger interface { diff --git a/testlib/events_observer_mock.go b/testlib/events_observer_mock.go new file mode 100644 index 0000000..e3d8979 --- /dev/null +++ b/testlib/events_observer_mock.go @@ -0,0 +1,7 @@ +package testlib + +import "github.com/stretchr/testify/mock" + +type EventsObserverMock struct { + mock.Mock +}