ScuroNeko/mtg
1
0
Fork 0
mirror of https://github.com/ScuroNeko/mtg.git synced 2026-08-31 17:24:02 +03:00
Code Issues Packages Projects Releases Wiki Activity

Add EventIPBlocklisted

Browse Source
This commit is contained in:
9seconds
2021-03-17 21:24:24 +03:00
parent 23519913f2
commit 2408f1530f
15 changed files with 128 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
events/event_stream.go
+2
View File
@@ -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
+23
View File
@@ -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
+1
View File
@@ -6,6 +6,7 @@ type Observer interface {
EventStart(mtglib.EventStart)
EventFinish(mtglib.EventFinish)
EventConcurrencyLimited(mtglib.EventConcurrencyLimited)
EventIPBlocklisted(mtglib.EventIPBlocklisted)
Shutdown()
}
events/init_test.go
+4
View File
@@ -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
+15
View File
@@ -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
+1
View File
@@ -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
+6
View File
@@ -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
+9
View File
@@ -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
+4
View File
@@ -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
+10
View File
@@ -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
+2 -1
View File
@@ -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
+16 -1
View File
@@ -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
+13
View File
@@ -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
+12
View File
@@ -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
+10
View File
@@ -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{})