Add EventIPBlocklisted

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)
 			}

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()

events/init.go

@@ -6,6 +6,7 @@ type Observer interface {
 	EventStart(mtglib.EventStart)
 	EventFinish(mtglib.EventFinish)
 	EventConcurrencyLimited(mtglib.EventConcurrencyLimited)
+	EventIPBlocklisted(mtglib.EventIPBlocklisted)
 
 	Shutdown()
 }

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()
 }

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()

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 {

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)
 			}
 		})
 	}

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 ""
+}

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{})

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 {

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"
 )

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
 }

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{})

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))

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{})