From 2408f1530fc8e026796d929dbcad63adebe77079 Mon Sep 17 00:00:00 2001 From: 9seconds Date: Wed, 17 Mar 2021 21:24:24 +0300 Subject: [PATCH] Add EventIPBlocklisted --- events/event_stream.go | 2 ++ events/event_stream_test.go | 23 +++++++++++++++++++++++ events/init.go | 1 + events/init_test.go | 4 ++++ events/multi_observer.go | 15 +++++++++++++++ events/noop.go | 1 + events/noop_test.go | 6 ++++++ mtglib/events.go | 9 +++++++++ mtglib/events_test.go | 4 ++++ mtglib/proxy.go | 10 ++++++++++ stats/init.go | 3 ++- stats/prometheus.go | 17 ++++++++++++++++- stats/prometheus_test.go | 13 +++++++++++++ stats/statsd.go | 12 ++++++++++++ stats/statsd_test.go | 10 ++++++++++ 15 files changed, 128 insertions(+), 2 deletions(-) diff --git a/events/event_stream.go b/events/event_stream.go index be860f6..5037fb6 100644 --- a/events/event_stream.go +++ b/events/event_stream.go @@ -73,6 +73,8 @@ func eventStreamProcessor(ctx context.Context, eventChan <-chan mtglib.Event, ob observer.EventStart(typedEvt) case mtglib.EventFinish: observer.EventFinish(typedEvt) + case mtglib.EventIPBlocklisted: + observer.EventIPBlocklisted(typedEvt) case mtglib.EventConcurrencyLimited: observer.EventConcurrencyLimited(typedEvt) } diff --git a/events/event_stream_test.go b/events/event_stream_test.go index b86b68b..ec708a0 100644 --- a/events/event_stream_test.go +++ b/events/event_stream_test.go @@ -106,6 +106,29 @@ func (suite *EventStreamTestSuite) TestEventConcurrencyLimitedOk() { time.Sleep(100 * time.Millisecond) } +func (suite *EventStreamTestSuite) TestEventIPBlocklistedOk() { + evt := mtglib.EventIPBlocklisted{ + CreatedAt: time.Now(), + RemoteIP: net.ParseIP("10.0.0.10"), + } + + for _, v := range []*ObserverMock{suite.observerMock1, suite.observerMock2} { + v. + On("EventIPBlocklisted", mock.Anything). + Once(). + Run(func(args mock.Arguments) { + caught := args.Get(0).(mtglib.EventIPBlocklisted) + + suite.Equal(evt.CreatedAt, caught.CreatedAt) + suite.Equal(evt.StreamID(), caught.StreamID()) + suite.Equal(evt.RemoteIP.String(), caught.RemoteIP.String()) + }) + } + + suite.stream.Send(suite.ctx, evt) + time.Sleep(100 * time.Millisecond) +} + func (suite *EventStreamTestSuite) TearDownTest() { suite.stream.Shutdown() suite.ctxCancel() diff --git a/events/init.go b/events/init.go index 0e348f5..6924071 100644 --- a/events/init.go +++ b/events/init.go @@ -6,6 +6,7 @@ type Observer interface { EventStart(mtglib.EventStart) EventFinish(mtglib.EventFinish) EventConcurrencyLimited(mtglib.EventConcurrencyLimited) + EventIPBlocklisted(mtglib.EventIPBlocklisted) Shutdown() } diff --git a/events/init_test.go b/events/init_test.go index e75e5e9..faa4063 100644 --- a/events/init_test.go +++ b/events/init_test.go @@ -21,6 +21,10 @@ func (o *ObserverMock) EventConcurrencyLimited(evt mtglib.EventConcurrencyLimite o.Called(evt) } +func (o *ObserverMock) EventIPBlocklisted(evt mtglib.EventIPBlocklisted) { + o.Called(evt) +} + func (o *ObserverMock) Shutdown() { o.Called() } diff --git a/events/multi_observer.go b/events/multi_observer.go index 698d87c..d97d802 100644 --- a/events/multi_observer.go +++ b/events/multi_observer.go @@ -55,6 +55,21 @@ func (m multiObserver) EventConcurrencyLimited(evt mtglib.EventConcurrencyLimite wg.Wait() } +func (m multiObserver) EventIPBlocklisted(evt mtglib.EventIPBlocklisted) { + wg := &sync.WaitGroup{} + wg.Add(len(m.observers)) + + for _, v := range m.observers { + go func(obs Observer) { + defer wg.Done() + + obs.EventIPBlocklisted(evt) + }(v) + } + + wg.Wait() +} + func (m multiObserver) Shutdown() { for _, v := range m.observers { v.Shutdown() diff --git a/events/noop.go b/events/noop.go index 7eb4041..7acbb88 100644 --- a/events/noop.go +++ b/events/noop.go @@ -20,6 +20,7 @@ type noopObserver struct{} func (n noopObserver) EventStart(_ mtglib.EventStart) {} func (n noopObserver) EventFinish(_ mtglib.EventFinish) {} func (n noopObserver) EventConcurrencyLimited(_ mtglib.EventConcurrencyLimited) {} +func (n noopObserver) EventIPBlocklisted(_ mtglib.EventIPBlocklisted) {} func (n noopObserver) Shutdown() {} func NewNoopObserver() Observer { diff --git a/events/noop_test.go b/events/noop_test.go index b5120f3..d3a3109 100644 --- a/events/noop_test.go +++ b/events/noop_test.go @@ -30,6 +30,10 @@ func (suite *NoopTestSuite) SetupSuite() { ConnID: "connID", }, "concurrency-limited": mtglib.EventConcurrencyLimited{}, + "ip-blacklisted": mtglib.EventIPBlocklisted{ + RemoteIP: net.ParseIP("10.0.0.10"), + CreatedAt: time.Now(), + }, } suite.ctx = context.Background() } @@ -62,6 +66,8 @@ func (suite *NoopTestSuite) TestObserver() { observer.EventFinish(typedEvt) case mtglib.EventConcurrencyLimited: observer.EventConcurrencyLimited(typedEvt) + case mtglib.EventIPBlocklisted: + observer.EventIPBlocklisted(typedEvt) } }) } diff --git a/mtglib/events.go b/mtglib/events.go index e7ea557..0e9a2d1 100644 --- a/mtglib/events.go +++ b/mtglib/events.go @@ -31,3 +31,12 @@ type EventConcurrencyLimited struct { func (e EventConcurrencyLimited) StreamID() string { return "" } + +type EventIPBlocklisted struct { + CreatedAt time.Time + RemoteIP net.IP +} + +func (e EventIPBlocklisted) StreamID() string { + return "" +} diff --git a/mtglib/events_test.go b/mtglib/events_test.go index ed08aa5..f0e3975 100644 --- a/mtglib/events_test.go +++ b/mtglib/events_test.go @@ -36,6 +36,10 @@ func (suite *EventsTestSuite) TestEventConcurrencyLimited() { suite.Empty(mtglib.EventConcurrencyLimited{}.StreamID()) } +func (suite *EventsTestSuite) TestEventIPBlocklisted() { + suite.Empty(mtglib.EventIPBlocklisted{}.StreamID()) +} + func TestEvents(t *testing.T) { t.Parallel() suite.Run(t, &EventsTestSuite{}) diff --git a/mtglib/proxy.go b/mtglib/proxy.go index e603db1..6c68415 100644 --- a/mtglib/proxy.go +++ b/mtglib/proxy.go @@ -52,6 +52,16 @@ func (p *Proxy) Serve(listener net.Listener) error { return fmt.Errorf("cannot accept a new connection: %w", err) } + if addr := conn.RemoteAddr().(*net.TCPAddr).IP; p.ipBlocklist.Contains(addr) { + conn.Close() + p.eventStream.Send(p.ctx, EventIPBlocklisted{ + CreatedAt: time.Now(), + RemoteIP: addr, + }) + + continue + } + err = p.workerPool.Invoke(conn) switch { diff --git a/stats/init.go b/stats/init.go index 7210888..063eb32 100644 --- a/stats/init.go +++ b/stats/init.go @@ -4,9 +4,10 @@ const ( MetricActiveConnection = "active_connections" MetricSessionDuration = "session_duration" MetricConcurrencyLimited = "concurrency_limited" + MetricIPBlocklisted = "ip_blocklisted" TagIPType = "ip_type" TagIPTypeIPv4 = "ipv4" - TagIPTypeIPv6 = "ipv4" + TagIPTypeIPv6 = "ipv6" ) diff --git a/stats/prometheus.go b/stats/prometheus.go index 69b23e3..be4c676 100644 --- a/stats/prometheus.go +++ b/stats/prometheus.go @@ -41,10 +41,18 @@ func (p prometheusProcessor) EventFinish(evt mtglib.EventFinish) { p.factory.metricSessionDuration.Observe(float64(duration) / float64(time.Second)) } -func (p prometheusProcessor) EventConcurrencyLimited(evt mtglib.EventConcurrencyLimited) { +func (p prometheusProcessor) EventConcurrencyLimited(_ mtglib.EventConcurrencyLimited) { p.factory.metricConcurrencyLimited.Inc() } +func (p prometheusProcessor) EventIPBlocklisted(evt mtglib.EventIPBlocklisted) { + if evt.RemoteIP.To4() == nil { + p.factory.metricIPBlocklisted.WithLabelValues(TagIPTypeIPv6).Inc() + } else { + p.factory.metricIPBlocklisted.WithLabelValues(TagIPTypeIPv4).Inc() + } +} + func (p prometheusProcessor) Shutdown() { p.streams = make(map[string]*streamInfo) } @@ -53,6 +61,7 @@ type PrometheusFactory struct { httpServer *http.Server metricActiveConnections *prometheus.GaugeVec + metricIPBlocklisted *prometheus.CounterVec metricConcurrencyLimited prometheus.Counter metricSessionDuration prometheus.Histogram } @@ -113,11 +122,17 @@ func NewPrometheus(metricPrefix, httpPath string) *PrometheusFactory { Name: MetricConcurrencyLimited, Help: "A number of sessions that were rejected by concurrency limiter.", }), + metricIPBlocklisted: prometheus.NewCounterVec(prometheus.CounterOpts{ + Namespace: metricPrefix, + Name: MetricIPBlocklisted, + Help: "A number of rejected sessions due to ip blocklisting", + }, []string{TagIPType}), } registry.MustRegister(factory.metricActiveConnections) registry.MustRegister(factory.metricSessionDuration) registry.MustRegister(factory.metricConcurrencyLimited) + registry.MustRegister(factory.metricIPBlocklisted) return factory } diff --git a/stats/prometheus_test.go b/stats/prometheus_test.go index 319a6d1..4b2175a 100644 --- a/stats/prometheus_test.go +++ b/stats/prometheus_test.go @@ -91,6 +91,19 @@ func (suite *PrometheusTestSuite) TestEventConcurrencyLimited() { suite.Contains(data, `mtg_concurrency_limited 1`) } +func (suite *PrometheusTestSuite) TestEventIPBlocklisted() { + suite.prometheus.EventIPBlocklisted(mtglib.EventIPBlocklisted{ + CreatedAt: time.Now(), + RemoteIP: net.ParseIP("2001:db8::68"), + }) + + time.Sleep(100 * time.Millisecond) + + data, err := suite.Get() + suite.NoError(err) + suite.Contains(data, `mtg_ip_blocklisted{ip_type="ipv6"} 1`) +} + func TestPrometheus(t *testing.T) { t.Parallel() suite.Run(t, &PrometheusTestSuite{}) diff --git a/stats/statsd.go b/stats/statsd.go index be04c51..69c8b37 100644 --- a/stats/statsd.go +++ b/stats/statsd.go @@ -49,6 +49,18 @@ func (s statsdProcessor) EventConcurrencyLimited(_ mtglib.EventConcurrencyLimite s.client.Incr(MetricConcurrencyLimited, 1) } +func (s statsdProcessor) EventIPBlocklisted(evt mtglib.EventIPBlocklisted) { + var tag statsd.Tag + + if evt.RemoteIP.To4() == nil { + tag = statsd.StringTag(TagIPType, TagIPTypeIPv6) + } else { + tag = statsd.StringTag(TagIPType, TagIPTypeIPv4) + } + + s.client.Incr(MetricIPBlocklisted, 1, tag) +} + func (s statsdProcessor) Shutdown() { now := time.Now() events := make([]mtglib.EventFinish, 0, len(s.streams)) diff --git a/stats/statsd_test.go b/stats/statsd_test.go index 8732463..1babb5b 100644 --- a/stats/statsd_test.go +++ b/stats/statsd_test.go @@ -121,6 +121,16 @@ func (suite *StatsdTestSuite) TestEventConcurrencyLimited() { suite.Equal("mtg.concurrency_limited:1|c", suite.statsdServer.String()) } +func (suite *StatsdTestSuite) TestEventIPBlocklisted() { + suite.statsd.EventIPBlocklisted(mtglib.EventIPBlocklisted{ + CreatedAt: time.Now(), + RemoteIP: net.ParseIP("10.0.0.10"), + }) + + time.Sleep(2 * statsd.DefaultFlushInterval) + suite.Equal("mtg.ip_blocklisted:1|c|#ip_type:ipv4", suite.statsdServer.String()) +} + func TestStatsd(t *testing.T) { t.Parallel() suite.Run(t, &StatsdTestSuite{})