diff --git a/events/event_stream.go b/events/event_stream.go index 5037fb6..1726c81 100644 --- a/events/event_stream.go +++ b/events/event_stream.go @@ -71,6 +71,10 @@ func eventStreamProcessor(ctx context.Context, eventChan <-chan mtglib.Event, ob switch typedEvt := evt.(type) { case mtglib.EventStart: observer.EventStart(typedEvt) + case mtglib.EventConnectedToDC: + observer.EventConnectedToDC(typedEvt) + case mtglib.EventTraffic: + observer.EventTraffic(typedEvt) case mtglib.EventFinish: observer.EventFinish(typedEvt) case mtglib.EventIPBlocklisted: diff --git a/events/event_stream_test.go b/events/event_stream_test.go index ec708a0..90dee03 100644 --- a/events/event_stream_test.go +++ b/events/event_stream_test.go @@ -38,7 +38,7 @@ func (suite *EventStreamTestSuite) SetupTest() { suite.stream = events.NewEventStream(factories) } -func (suite *EventStreamTestSuite) TestEventStartOk() { +func (suite *EventStreamTestSuite) TestEventStart() { evt := mtglib.EventStart{ CreatedAt: time.Now(), ConnID: "connID", @@ -63,7 +63,61 @@ func (suite *EventStreamTestSuite) TestEventStartOk() { time.Sleep(100 * time.Millisecond) } -func (suite *EventStreamTestSuite) TestEventFinishOk() { +func (suite *EventStreamTestSuite) TestEventConnectedToDC() { + evt := mtglib.EventConnectedToDC{ + CreatedAt: time.Now(), + ConnID: "connID", + RemoteIP: net.ParseIP("10.0.0.1"), + DC: 3, + } + + for _, v := range []*ObserverMock{suite.observerMock1, suite.observerMock2} { + v. + On("EventConnectedToDC", mock.Anything). + Once(). + Run(func(args mock.Arguments) { + caught := args.Get(0).(mtglib.EventConnectedToDC) + + suite.Equal(evt.CreatedAt, caught.CreatedAt) + suite.Equal(evt.ConnID, caught.ConnID) + suite.Equal(evt.RemoteIP.String(), caught.RemoteIP.String()) + suite.Equal(evt.StreamID(), caught.StreamID()) + suite.Equal(evt.DC, caught.DC) + }) + } + + suite.stream.Send(suite.ctx, evt) + time.Sleep(100 * time.Millisecond) +} + +func (suite *EventStreamTestSuite) TestEventTraffic() { + evt := mtglib.EventTraffic{ + CreatedAt: time.Now(), + ConnID: "connID", + Traffic: 1024, + IsRead: true, + } + + for _, v := range []*ObserverMock{suite.observerMock1, suite.observerMock2} { + v. + On("EventTraffic", mock.Anything). + Once(). + Run(func(args mock.Arguments) { + caught := args.Get(0).(mtglib.EventTraffic) + + suite.Equal(evt.CreatedAt, caught.CreatedAt) + suite.Equal(evt.ConnID, caught.ConnID) + suite.Equal(evt.StreamID(), caught.StreamID()) + suite.Equal(evt.Traffic, caught.Traffic) + suite.Equal(evt.IsRead, caught.IsRead) + }) + } + + suite.stream.Send(suite.ctx, evt) + time.Sleep(100 * time.Millisecond) +} + +func (suite *EventStreamTestSuite) TestEventFinish() { evt := mtglib.EventFinish{ CreatedAt: time.Now(), ConnID: "connID", @@ -86,7 +140,7 @@ func (suite *EventStreamTestSuite) TestEventFinishOk() { time.Sleep(100 * time.Millisecond) } -func (suite *EventStreamTestSuite) TestEventConcurrencyLimitedOk() { +func (suite *EventStreamTestSuite) TestEventConcurrencyLimited() { evt := mtglib.EventConcurrencyLimited{ CreatedAt: time.Now(), } @@ -106,7 +160,7 @@ func (suite *EventStreamTestSuite) TestEventConcurrencyLimitedOk() { time.Sleep(100 * time.Millisecond) } -func (suite *EventStreamTestSuite) TestEventIPBlocklistedOk() { +func (suite *EventStreamTestSuite) TestEventIPBlocklisted() { evt := mtglib.EventIPBlocklisted{ CreatedAt: time.Now(), RemoteIP: net.ParseIP("10.0.0.10"), diff --git a/events/init.go b/events/init.go index 6924071..4720a3b 100644 --- a/events/init.go +++ b/events/init.go @@ -5,6 +5,8 @@ import "github.com/9seconds/mtg/v2/mtglib" type Observer interface { EventStart(mtglib.EventStart) EventFinish(mtglib.EventFinish) + EventConnectedToDC(mtglib.EventConnectedToDC) + EventTraffic(mtglib.EventTraffic) EventConcurrencyLimited(mtglib.EventConcurrencyLimited) EventIPBlocklisted(mtglib.EventIPBlocklisted) diff --git a/events/init_test.go b/events/init_test.go index faa4063..58e3635 100644 --- a/events/init_test.go +++ b/events/init_test.go @@ -13,6 +13,14 @@ func (o *ObserverMock) EventStart(evt mtglib.EventStart) { o.Called(evt) } +func (o *ObserverMock) EventConnectedToDC(evt mtglib.EventConnectedToDC) { + o.Called(evt) +} + +func (o *ObserverMock) EventTraffic(evt mtglib.EventTraffic) { + o.Called(evt) +} + func (o *ObserverMock) EventFinish(evt mtglib.EventFinish) { o.Called(evt) } diff --git a/events/multi_observer.go b/events/multi_observer.go index d97d802..1dec489 100644 --- a/events/multi_observer.go +++ b/events/multi_observer.go @@ -25,6 +25,36 @@ func (m multiObserver) EventStart(evt mtglib.EventStart) { wg.Wait() } +func (m multiObserver) EventConnectedToDC(evt mtglib.EventConnectedToDC) { + wg := &sync.WaitGroup{} + wg.Add(len(m.observers)) + + for _, v := range m.observers { + go func(obs Observer) { + defer wg.Done() + + obs.EventConnectedToDC(evt) + }(v) + } + + wg.Wait() +} + +func (m multiObserver) EventTraffic(evt mtglib.EventTraffic) { + wg := &sync.WaitGroup{} + wg.Add(len(m.observers)) + + for _, v := range m.observers { + go func(obs Observer) { + defer wg.Done() + + obs.EventTraffic(evt) + }(v) + } + + wg.Wait() +} + func (m multiObserver) EventFinish(evt mtglib.EventFinish) { wg := &sync.WaitGroup{} wg.Add(len(m.observers)) diff --git a/events/noop.go b/events/noop.go index 7acbb88..8dcbf57 100644 --- a/events/noop.go +++ b/events/noop.go @@ -18,6 +18,8 @@ func NewNoopStream() mtglib.EventStream { type noopObserver struct{} func (n noopObserver) EventStart(_ mtglib.EventStart) {} +func (n noopObserver) EventConnectedToDC(_ mtglib.EventConnectedToDC) {} +func (n noopObserver) EventTraffic(_ mtglib.EventTraffic) {} func (n noopObserver) EventFinish(_ mtglib.EventFinish) {} func (n noopObserver) EventConcurrencyLimited(_ mtglib.EventConcurrencyLimited) {} func (n noopObserver) EventIPBlocklisted(_ mtglib.EventIPBlocklisted) {} diff --git a/events/noop_test.go b/events/noop_test.go index d3a3109..2621e38 100644 --- a/events/noop_test.go +++ b/events/noop_test.go @@ -25,6 +25,18 @@ func (suite *NoopTestSuite) SetupSuite() { ConnID: "connID", RemoteIP: net.ParseIP("127.0.0.1"), }, + "connected-to-dc": mtglib.EventConnectedToDC{ + CreatedAt: time.Now(), + ConnID: "connID", + RemoteIP: net.ParseIP("127.1.0.1"), + DC: 2, + }, + "traffic": mtglib.EventTraffic{ + CreatedAt: time.Now(), + ConnID: "connID", + Traffic: 1000, + IsRead: true, + }, "finish": mtglib.EventFinish{ CreatedAt: time.Now(), ConnID: "connID", @@ -62,6 +74,8 @@ func (suite *NoopTestSuite) TestObserver() { switch typedEvt := value.(type) { case mtglib.EventStart: observer.EventStart(typedEvt) + case mtglib.EventConnectedToDC: + observer.EventConnectedToDC(typedEvt) case mtglib.EventFinish: observer.EventFinish(typedEvt) case mtglib.EventConcurrencyLimited: diff --git a/mtglib/conns.go b/mtglib/conns.go new file mode 100644 index 0000000..685c7db --- /dev/null +++ b/mtglib/conns.go @@ -0,0 +1,91 @@ +package mtglib + +import ( + "context" + "fmt" + "net" + "time" +) + +type connStandard struct { + conn net.Conn + idleTimeout time.Duration +} + +func (c connStandard) Read(b []byte) (int, error) { + if err := c.conn.SetReadDeadline(time.Now().Add(c.idleTimeout)); err != nil { + return 0, fmt.Errorf("cannot set read deadline: %w", err) + } + + return c.conn.Read(b) +} + +func (c connStandard) Write(b []byte) (int, error) { + if err := c.conn.SetWriteDeadline(time.Now().Add(c.idleTimeout)); err != nil { + return 0, fmt.Errorf("cannot set write deadline: %w", err) + } + + return c.conn.Write(b) +} + +func (c connStandard) Close() error { + return c.conn.Close() +} + +func (c connStandard) LocalAddr() net.Addr { + return c.conn.LocalAddr() +} + +func (c connStandard) RemoteAddr() net.Addr { + return c.conn.RemoteAddr() +} + +func (c connStandard) SetDeadline(t time.Time) error { + return c.conn.SetDeadline(t) +} + +func (c connStandard) SetReadDeadline(t time.Time) error { + return c.conn.SetReadDeadline(t) +} + +func (c connStandard) SetWriteDeadline(t time.Time) error { + return c.conn.SetWriteDeadline(t) +} + +type connEventTraffic struct { + net.Conn + + connID string + stream EventStream + ctx context.Context +} + +func (c connEventTraffic) Read(b []byte) (int, error) { + n, err := c.Conn.Read(b) + + if n > 0 { + c.stream.Send(c.ctx, EventTraffic{ + CreatedAt: time.Now(), + ConnID: c.connID, + Traffic: uint(n), + IsRead: true, + }) + } + + return n, err // nolint: wrapcheck +} + +func (c connEventTraffic) Write(b []byte) (int, error) { + n, err := c.Conn.Write(b) + + if n > 0 { + c.stream.Send(c.ctx, EventTraffic{ + CreatedAt: time.Now(), + ConnID: c.connID, + Traffic: uint(n), + IsRead: false, + }) + } + + return n, err // nolint: wrapcheck +} diff --git a/mtglib/events.go b/mtglib/events.go index 0e9a2d1..0e49e48 100644 --- a/mtglib/events.go +++ b/mtglib/events.go @@ -15,6 +15,28 @@ func (e EventStart) StreamID() string { return e.ConnID } +type EventConnectedToDC struct { + CreatedAt time.Time + ConnID string + RemoteIP net.IP + DC int +} + +func (e EventConnectedToDC) StreamID() string { + return e.ConnID +} + +type EventTraffic struct { + CreatedAt time.Time + ConnID string + Traffic uint + IsRead bool +} + +func (e EventTraffic) StreamID() string { + return e.ConnID +} + type EventFinish struct { CreatedAt time.Time ConnID string diff --git a/mtglib/internal/obfuscated2/conn.go b/mtglib/internal/obfuscated2/conn.go index dbd69ed..7e25d47 100644 --- a/mtglib/internal/obfuscated2/conn.go +++ b/mtglib/internal/obfuscated2/conn.go @@ -14,7 +14,7 @@ type Conn struct { writeBuf []byte } -func (c *Conn) Read(p []byte) (int, error) { +func (c Conn) Read(p []byte) (int, error) { n, err := c.Conn.Read(p) if err != nil { return n, err // nolint: wrapcheck @@ -25,7 +25,7 @@ func (c *Conn) Read(p []byte) (int, error) { return n, nil } -func (c *Conn) Write(p []byte) (int, error) { +func (c Conn) Write(p []byte) (int, error) { c.writeBuf = append(c.writeBuf[:0], p...) c.Encryptor.XORKeyStream(c.writeBuf, c.writeBuf) diff --git a/mtglib/proxy.go b/mtglib/proxy.go index 312a1bf..9cedf31 100644 --- a/mtglib/proxy.go +++ b/mtglib/proxy.go @@ -9,6 +9,7 @@ import ( "time" "github.com/9seconds/mtg/v2/mtglib/internal/obfuscated2" + "github.com/9seconds/mtg/v2/mtglib/internal/telegram" "github.com/panjf2000/ants/v2" ) @@ -16,10 +17,12 @@ type Proxy struct { ctx context.Context ctxCancel context.CancelFunc streamWaitGroup sync.WaitGroup - workerPool *ants.PoolWithFunc + + idleTimeout time.Duration + workerPool *ants.PoolWithFunc + telegram *telegram.Telegram secret Secret - network Network antiReplayCache AntiReplayCache ipBlocklist IPBlocklist eventStream EventStream @@ -55,6 +58,12 @@ func (p *Proxy) ServeConn(conn net.Conn) { return } + + if err := p.doTelegramCall(ctx); err != nil { + p.logger.WarningError("cannot dial to telegram", err) + + return + } } func (p *Proxy) Serve(listener net.Listener) error { @@ -102,16 +111,45 @@ func (p *Proxy) doObfuscated2Handshake(ctx *streamContext) error { ctx.dc = dc ctx.logger = ctx.logger.BindInt("dc", dc) - ctx.clientConn = &obfuscated2.Conn{ - Conn: ctx.clientConn, - Encryptor: encryptor, - Decryptor: decryptor, + ctx.clientConn = connStandard{ + conn: obfuscated2.Conn{ + Conn: ctx.clientConn, + Encryptor: encryptor, + Decryptor: decryptor, + }, + idleTimeout: p.idleTimeout, } return nil } -func NewProxy(opts ProxyOpts) (*Proxy, error) { +func (p *Proxy) doTelegramCall(ctx *streamContext) error { + conn, err := p.telegram.Dial(ctx, ctx.dc) + if err != nil { + return fmt.Errorf("cannot dial to Telegram: %w", err) + } + + ctx.telegramConn = connEventTraffic{ + Conn: connStandard{ + conn: conn, + idleTimeout: p.idleTimeout, + }, + connID: ctx.connID, + stream: p.eventStream, + ctx: ctx, + } + + p.eventStream.Send(ctx, EventConnectedToDC{ + CreatedAt: time.Now(), + ConnID: ctx.connID, + RemoteIP: conn.RemoteAddr().(*net.TCPAddr).IP, + DC: ctx.dc, + }) + + return nil +} + +func NewProxy(opts ProxyOpts) (*Proxy, error) { // nolint: cyclop switch { case opts.Network == nil: return nil, ErrNetworkIsNotDefined @@ -127,21 +165,32 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) { return nil, ErrSecretInvalid } + tg, err := telegram.New(opts.Network, opts.PreferIP) + if err != nil { + return nil, fmt.Errorf("cannot build telegram dialer: %w", err) + } + concurrency := opts.Concurrency if concurrency == 0 { concurrency = DefaultConcurrency } + idleTimeout := opts.IdleTimeout + if idleTimeout < 1 { + idleTimeout = DefaultIdleTimeout + } + ctx, cancel := context.WithCancel(context.Background()) proxy := &Proxy{ ctx: ctx, ctxCancel: cancel, secret: opts.Secret, - network: opts.Network, antiReplayCache: opts.AntiReplayCache, ipBlocklist: opts.IPBlocklist, eventStream: opts.EventStream, logger: opts.Logger.Named("proxy"), + idleTimeout: idleTimeout, + telegram: tg, } pool, err := ants.NewPoolWithFunc(int(concurrency), func(arg interface{}) { diff --git a/mtglib/stream_context.go b/mtglib/stream_context.go index 1c75482..2ab91a0 100644 --- a/mtglib/stream_context.go +++ b/mtglib/stream_context.go @@ -9,12 +9,13 @@ import ( ) type streamContext struct { - ctx context.Context - ctxCancel context.CancelFunc - clientConn net.Conn - connID string - dc int - logger Logger + ctx context.Context + ctxCancel context.CancelFunc + clientConn net.Conn + telegramConn net.Conn + connID string + dc int + logger Logger } func (s *streamContext) Deadline() (time.Time, bool) { @@ -35,7 +36,14 @@ func (s *streamContext) Value(key interface{}) interface{} { func (s *streamContext) Close() { s.ctxCancel() - s.clientConn.Close() + + if s.clientConn != nil { + s.clientConn.Close() + } + + if s.telegramConn != nil { + s.telegramConn.Close() + } } func (s *streamContext) ClientIP() net.IP { diff --git a/stats/init.go b/stats/init.go index 7d095c5..4bf1aa8 100644 --- a/stats/init.go +++ b/stats/init.go @@ -6,13 +6,21 @@ const ( DefaultStatsdMetricPrefix = DefaultMetricPrefix + "." DefaultStatsdTagFormat = "datadog" - MetricActiveConnection = "active_connections" - MetricSessionDuration = "session_duration" - MetricConcurrencyLimited = "concurrency_limited" - MetricIPBlocklisted = "ip_blocklisted" + MetricClientConnections = "client_connections" + MetricTelegramConnections = "telegram_connections" + MetricTraffic = "traffic" + MetricSessionDuration = "session_duration" + MetricSessionTraffic = "session_traffic" + MetricConcurrencyLimited = "concurrency_limited" + MetricIPBlocklisted = "ip_blocklisted" - TagIPType = "ip_type" + TagIPType = "ip_type" + TagTelegramIP = "ip" + TagDC = "dc" + TagDirection = "direction" - TagIPTypeIPv4 = "ipv4" - TagIPTypeIPv6 = "ipv6" + TagIPTypeIPv4 = "ipv4" + TagIPTypeIPv6 = "ipv6" + TagDirectionTelegram = "telegram" + TagDirectionClient = "client" ) diff --git a/stats/prometheus.go b/stats/prometheus.go index be4c676..4c9b0a3 100644 --- a/stats/prometheus.go +++ b/stats/prometheus.go @@ -4,6 +4,7 @@ import ( "context" "net" "net/http" + "strconv" "time" "github.com/9seconds/mtg/v2/events" @@ -24,7 +25,47 @@ func (p prometheusProcessor) EventStart(evt mtglib.EventStart) { } p.streams[evt.StreamID()] = sInfo - p.factory.metricActiveConnections.WithLabelValues(sInfo.IPType()).Inc() + p.factory.metricClientConnections.WithLabelValues(sInfo.GetClientIPType()).Inc() +} + +func (p prometheusProcessor) EventConnectedToDC(evt mtglib.EventConnectedToDC) { + sInfo, ok := p.streams[evt.StreamID()] + if !ok { + return + } + + sInfo.remoteIP = evt.RemoteIP + sInfo.dc = evt.DC + + p.factory.metricTelegramConnections.WithLabelValues( + sInfo.GetRemoteIPType(), + sInfo.remoteIP.String(), + strconv.Itoa(sInfo.dc)).Inc() +} + +func (p prometheusProcessor) EventTraffic(evt mtglib.EventTraffic) { + sInfo, ok := p.streams[evt.StreamID()] + if !ok { + return + } + + labels := []string{ + sInfo.GetRemoteIPType(), + sInfo.remoteIP.String(), + strconv.Itoa(sInfo.dc), + } + + if evt.IsRead { + sInfo.bytesRecvFromTelegram += evt.Traffic + + labels = append(labels, TagDirectionClient) + } else { + sInfo.bytesSentToTelegram += evt.Traffic + + labels = append(labels, TagDirectionTelegram) + } + + p.factory.metricTraffic.WithLabelValues(labels...).Add(float64(evt.Traffic)) } func (p prometheusProcessor) EventFinish(evt mtglib.EventFinish) { @@ -37,8 +78,30 @@ func (p prometheusProcessor) EventFinish(evt mtglib.EventFinish) { duration := evt.CreatedAt.Sub(sInfo.createdAt) - p.factory.metricActiveConnections.WithLabelValues(sInfo.IPType()).Dec() + p.factory.metricClientConnections.WithLabelValues(sInfo.GetRemoteIPType()).Dec() p.factory.metricSessionDuration.Observe(float64(duration) / float64(time.Second)) + + if sInfo.remoteIP == nil { + return + } + + labels := []string{ + sInfo.GetRemoteIPType(), + sInfo.remoteIP.String(), + strconv.Itoa(sInfo.dc), + } + + p.factory.metricTelegramConnections.WithLabelValues(labels...).Dec() + + labels = append(labels, TagDirectionClient) + p.factory.metricSessionTraffic. + WithLabelValues(labels...). + Observe(float64(sInfo.bytesRecvFromTelegram)) + + labels[3] = TagDirectionTelegram + p.factory.metricSessionTraffic. + WithLabelValues(labels...). + Observe(float64(sInfo.bytesSentToTelegram)) } func (p prometheusProcessor) EventConcurrencyLimited(_ mtglib.EventConcurrencyLimited) { @@ -60,10 +123,13 @@ func (p prometheusProcessor) Shutdown() { type PrometheusFactory struct { httpServer *http.Server - metricActiveConnections *prometheus.GaugeVec - metricIPBlocklisted *prometheus.CounterVec - metricConcurrencyLimited prometheus.Counter - metricSessionDuration prometheus.Histogram + metricClientConnections *prometheus.GaugeVec + metricTelegramConnections *prometheus.GaugeVec + metricTraffic *prometheus.CounterVec + metricIPBlocklisted *prometheus.CounterVec + metricSessionTraffic *prometheus.HistogramVec + metricConcurrencyLimited prometheus.Counter + metricSessionDuration prometheus.Histogram } func (p *PrometheusFactory) Make() events.Observer { @@ -81,7 +147,7 @@ func (p *PrometheusFactory) Close() error { return p.httpServer.Shutdown(context.Background()) } -func NewPrometheus(metricPrefix, httpPath string) *PrometheusFactory { +func NewPrometheus(metricPrefix, httpPath string) *PrometheusFactory { // nolint: funlen registry := prometheus.NewPedanticRegistry() httpHandler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{ EnableOpenMetrics: true, @@ -95,11 +161,16 @@ func NewPrometheus(metricPrefix, httpPath string) *PrometheusFactory { Handler: mux, }, - metricActiveConnections: prometheus.NewGaugeVec(prometheus.GaugeOpts{ + metricClientConnections: prometheus.NewGaugeVec(prometheus.GaugeOpts{ Namespace: metricPrefix, - Name: MetricActiveConnection, + Name: MetricClientConnections, Help: "A number of connections under active processing.", }, []string{TagIPType}), + metricTelegramConnections: prometheus.NewGaugeVec(prometheus.GaugeOpts{ + Namespace: metricPrefix, + Name: MetricTelegramConnections, + Help: "A number of connections to Telegram servers.", + }, []string{TagIPType, TagTelegramIP, TagDC}), metricSessionDuration: prometheus.NewHistogram(prometheus.HistogramOpts{ Namespace: metricPrefix, Name: MetricSessionDuration, @@ -117,6 +188,27 @@ func NewPrometheus(metricPrefix, httpPath string) *PrometheusFactory { 300, }, }), + metricSessionTraffic: prometheus.NewHistogramVec(prometheus.HistogramOpts{ + Namespace: metricPrefix, + Name: MetricSessionTraffic, + Help: "A traffic size which flew via proxy within a single session.", + Buckets: []float64{ // per 1mb + 1 * 1024 * 1024, + 2 * 1024 * 1024, + 3 * 1024 * 1024, + 4 * 1024 * 1024, + 5 * 1024 * 1024, + 6 * 1024 * 1024, + 7 * 1024 * 1024, + 8 * 1024 * 1024, + 9 * 1024 * 1024, + }, + }, []string{TagIPType, TagTelegramIP, TagDC, TagDirection}), + metricTraffic: prometheus.NewCounterVec(prometheus.CounterOpts{ + Namespace: metricPrefix, + Name: MetricTraffic, + Help: "Traffic which is sent through this proxy.", + }, []string{TagIPType, TagTelegramIP, TagDC, TagDirection}), metricConcurrencyLimited: prometheus.NewCounter(prometheus.CounterOpts{ Namespace: metricPrefix, Name: MetricConcurrencyLimited, @@ -129,7 +221,10 @@ func NewPrometheus(metricPrefix, httpPath string) *PrometheusFactory { }, []string{TagIPType}), } - registry.MustRegister(factory.metricActiveConnections) + registry.MustRegister(factory.metricClientConnections) + registry.MustRegister(factory.metricTelegramConnections) + registry.MustRegister(factory.metricTraffic) + registry.MustRegister(factory.metricSessionTraffic) registry.MustRegister(factory.metricSessionDuration) registry.MustRegister(factory.metricConcurrencyLimited) registry.MustRegister(factory.metricIPBlocklisted) diff --git a/stats/prometheus_test.go b/stats/prometheus_test.go index 00623cd..c1425b5 100644 --- a/stats/prometheus_test.go +++ b/stats/prometheus_test.go @@ -60,23 +60,60 @@ func (suite *PrometheusTestSuite) TestEventStartFinish() { ConnID: "connID", RemoteIP: net.ParseIP("10.0.0.10"), }) - time.Sleep(100 * time.Millisecond) data, err := suite.Get() suite.NoError(err) - suite.Contains(data, `mtg_active_connections{ip_type="ipv4"} 1`) + suite.Contains(data, `mtg_client_connections{ip_type="ipv4"} 1`) + + suite.prometheus.EventConnectedToDC(mtglib.EventConnectedToDC{ + CreatedAt: time.Now(), + ConnID: "connID", + RemoteIP: net.ParseIP("10.0.0.1"), + DC: 4, + }) + time.Sleep(100 * time.Millisecond) + + data, err = suite.Get() + suite.NoError(err) + suite.Contains(data, `mtg_telegram_connections{dc="4",ip="10.0.0.1",ip_type="ipv4"} 1`) + + suite.prometheus.EventTraffic(mtglib.EventTraffic{ + CreatedAt: time.Now(), + ConnID: "connID", + Traffic: 200, + IsRead: true, + }) + time.Sleep(100 * time.Millisecond) + + data, err = suite.Get() + suite.NoError(err) + suite.Contains(data, `mtg_traffic{dc="4",direction="client",ip="10.0.0.1",ip_type="ipv4"} 200`) + + suite.prometheus.EventTraffic(mtglib.EventTraffic{ + CreatedAt: time.Now(), + ConnID: "connID", + Traffic: 100, + IsRead: false, + }) + time.Sleep(100 * time.Millisecond) + + data, err = suite.Get() + suite.NoError(err) + suite.Contains(data, `mtg_traffic{dc="4",direction="telegram",ip="10.0.0.1",ip_type="ipv4"} 100`) suite.prometheus.EventFinish(mtglib.EventFinish{ CreatedAt: time.Now(), ConnID: "connID", }) - time.Sleep(100 * time.Millisecond) data, err = suite.Get() suite.NoError(err) - suite.Contains(data, `mtg_active_connections{ip_type="ipv4"} 0`) + suite.Contains(data, `mtg_client_connections{ip_type="ipv4"} 0`) + suite.Contains(data, `mtg_telegram_connections{dc="4",ip="10.0.0.1",ip_type="ipv4"} 0`) + suite.Contains(data, `mtg_traffic{dc="4",direction="client",ip="10.0.0.1",ip_type="ipv4"} 200`) + suite.Contains(data, `mtg_traffic{dc="4",direction="telegram",ip="10.0.0.1",ip_type="ipv4"} 100`) } func (suite *PrometheusTestSuite) TestEventConcurrencyLimited() { diff --git a/stats/statsd.go b/stats/statsd.go index f7c437c..beea0f8 100644 --- a/stats/statsd.go +++ b/stats/statsd.go @@ -22,9 +22,47 @@ func (s statsdProcessor) EventStart(evt mtglib.EventStart) { clientIP: evt.RemoteIP, } s.streams[evt.StreamID()] = sInfo - ipTypeTag := statsd.StringTag(TagIPType, sInfo.IPType()) - s.client.GaugeDelta(MetricActiveConnection, 1, ipTypeTag) + s.client.GaugeDelta(MetricClientConnections, + 1, + statsd.StringTag(TagIPType, sInfo.GetClientIPType())) +} + +func (s statsdProcessor) EventConnectedToDC(evt mtglib.EventConnectedToDC) { + sInfo, ok := s.streams[evt.StreamID()] + if !ok { + return + } + + sInfo.remoteIP = evt.RemoteIP + sInfo.dc = evt.DC + + s.client.GaugeDelta(MetricTelegramConnections, + 1, + statsd.StringTag(TagIPType, sInfo.GetRemoteIPType()), + statsd.StringTag(TagTelegramIP, sInfo.remoteIP.String()), + statsd.IntTag(TagDC, sInfo.dc)) +} + +func (s statsdProcessor) EventTraffic(evt mtglib.EventTraffic) { + sInfo, ok := s.streams[evt.StreamID()] + if !ok { + return + } + + tags := []statsd.Tag{ + statsd.StringTag(TagIPType, sInfo.GetRemoteIPType()), + statsd.StringTag(TagTelegramIP, sInfo.remoteIP.String()), + statsd.IntTag(TagDC, sInfo.dc), + } + + if evt.IsRead { + tags = append(tags, statsd.StringTag(TagDirection, TagDirectionClient)) + s.client.Incr(MetricTraffic, int64(evt.Traffic), tags...) + } else { + tags = append(tags, statsd.StringTag(TagDirection, TagDirectionTelegram)) + s.client.Incr(MetricTraffic, int64(evt.Traffic), tags...) + } } func (s statsdProcessor) EventFinish(evt mtglib.EventFinish) { @@ -35,11 +73,19 @@ func (s statsdProcessor) EventFinish(evt mtglib.EventFinish) { defer delete(s.streams, evt.StreamID()) - duration := evt.CreatedAt.Sub(sInfo.createdAt) - ipTypeTag := statsd.StringTag(TagIPType, sInfo.IPType()) + s.client.GaugeDelta(MetricClientConnections, + -1, + statsd.StringTag(TagIPType, sInfo.GetClientIPType())) + s.client.PrecisionTiming(MetricSessionDuration, + evt.CreatedAt.Sub(sInfo.createdAt)) - s.client.GaugeDelta(MetricActiveConnection, -1, ipTypeTag) - s.client.PrecisionTiming(MetricSessionDuration, duration) + if sInfo.remoteIP != nil { + s.client.GaugeDelta(MetricTelegramConnections, + -1, + statsd.StringTag(TagIPType, sInfo.GetRemoteIPType()), + statsd.StringTag(TagTelegramIP, sInfo.remoteIP.String()), + statsd.IntTag(TagDC, sInfo.dc)) + } } func (s statsdProcessor) EventConcurrencyLimited(_ mtglib.EventConcurrencyLimited) { diff --git a/stats/statsd_test.go b/stats/statsd_test.go index af9491f..9ac40bd 100644 --- a/stats/statsd_test.go +++ b/stats/statsd_test.go @@ -103,17 +103,49 @@ func (suite *StatsdTestSuite) TestEventStartFinish() { ConnID: "connID", RemoteIP: net.ParseIP("10.0.0.10"), }) - time.Sleep(statsdSleepTime) - suite.Equal("mtg.active_connections:+1|g|#ip_type:ipv4", suite.statsdServer.String()) + suite.Equal("mtg.client_connections:+1|g|#ip_type:ipv4", suite.statsdServer.String()) + + suite.statsd.EventConnectedToDC(mtglib.EventConnectedToDC{ + CreatedAt: time.Now(), + ConnID: "connID", + RemoteIP: net.ParseIP("10.1.0.10"), + DC: 2, + }) + time.Sleep(statsdSleepTime) + suite.Contains(suite.statsdServer.String(), + "mtg.telegram_connections:+1|g|#ip_type:ipv4,ip:10.1.0.10,dc:2") + + suite.statsd.EventTraffic(mtglib.EventTraffic{ + CreatedAt: time.Now(), + ConnID: "connID", + Traffic: 30, + IsRead: true, + }) + time.Sleep(statsdSleepTime) + suite.Contains(suite.statsdServer.String(), + "mtg.traffic:30|c|#ip_type:ipv4,ip:10.1.0.10,dc:2,direction:client") + + suite.statsd.EventTraffic(mtglib.EventTraffic{ + CreatedAt: time.Now(), + ConnID: "connID", + Traffic: 90, + IsRead: false, + }) + time.Sleep(statsdSleepTime) + suite.Contains(suite.statsdServer.String(), + "mtg.traffic:90|c|#ip_type:ipv4,ip:10.1.0.10,dc:2,direction:telegram") suite.statsd.EventFinish(mtglib.EventFinish{ CreatedAt: time.Now(), ConnID: "connID", }) - time.Sleep(statsdSleepTime) suite.Contains(suite.statsdServer.String(), "mtg.session_duration") + suite.Contains(suite.statsdServer.String(), + "mtg.telegram_connections:-1|g|#ip_type:ipv4,ip:10.1.0.10,dc:2") + suite.Contains(suite.statsdServer.String(), + "mtg.client_connections:-1|g|#ip_type:ipv4") } func (suite *StatsdTestSuite) TestEventConcurrencyLimited() { diff --git a/stats/stream_info.go b/stats/stream_info.go index fe51136..2827792 100644 --- a/stats/stream_info.go +++ b/stats/stream_info.go @@ -6,12 +6,24 @@ import ( ) type streamInfo struct { - createdAt time.Time - clientIP net.IP + createdAt time.Time + clientIP net.IP + remoteIP net.IP + dc int + bytesSentToTelegram uint + bytesRecvFromTelegram uint } -func (s *streamInfo) IPType() string { - if s.clientIP.To4() == nil { +func (s *streamInfo) GetClientIPType() string { + return s.getIPType(s.clientIP) +} + +func (s *streamInfo) GetRemoteIPType() string { + return s.getIPType(s.remoteIP) +} + +func (s *streamInfo) getIPType(ip net.IP) string { + if ip.To4() == nil { return TagIPTypeIPv6 }