Add EventTraffic

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

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 {