Add EventTraffic

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:

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

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)
 

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

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

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

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:

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

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

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)
 

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

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 {

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

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)

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.EventFinish(mtglib.EventFinish{
+	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() {

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(MetricActiveConnection, -1, ipTypeTag)
-	s.client.PrecisionTiming(MetricSessionDuration, duration)
+	s.client.GaugeDelta(MetricClientConnections,
+		-1,
+		statsd.StringTag(TagIPType, sInfo.GetClientIPType()))
+	s.client.PrecisionTiming(MetricSessionDuration,
+		evt.CreatedAt.Sub(sInfo.createdAt))
+
+	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) {

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.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.Equal("mtg.active_connections:+1|g|#ip_type:ipv4", suite.statsdServer.String())
+	suite.Contains(suite.statsdServer.String(),
+		"mtg.telegram_connections:+1|g|#ip_type:ipv4,ip:10.1.0.10,dc:2")
 
-	suite.statsd.EventFinish(mtglib.EventFinish{
+	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() {

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
 	}