Proxy is working in a simple mode now

mtglib/conns.go

@@ -26,7 +26,7 @@ func (c connTelegramTraffic) Read(b []byte) (int, error) {
 		})
 	}
 
-	return n, err
+	return n, err // nolint: wrapcheck
 }
 
 func (c connTelegramTraffic) Write(b []byte) (int, error) {
@@ -41,5 +41,5 @@ func (c connTelegramTraffic) Write(b []byte) (int, error) {
 		})
 	}
 
-	return n, err
+	return n, err // nolint: wrapcheck
 }

mtglib/internal/obfuscated2/client_handshake.go

@@ -23,16 +23,15 @@ func (c *clientHandhakeFrame) decryptor(secret []byte) cipher.Stream {
 }
 
 func (c *clientHandhakeFrame) encryptor(secret []byte) cipher.Stream {
-	arr := clientHandhakeFrame{}
+	invertedHandshake := c.invert()
-	invertByteSlices(arr.data[:], c.data[:])
 
 	hasher := acquireSha256Hasher()
 	defer releaseSha256Hasher(hasher)
 
-	hasher.Write(arr.key()) // nolint: errcheck
+	hasher.Write(invertedHandshake.key()) // nolint: errcheck
-	hasher.Write(secret)    // nolint: errcheck
+	hasher.Write(secret)                  // nolint: errcheck
 
-	return makeAesCtr(hasher.Sum(nil), arr.iv())
+	return makeAesCtr(hasher.Sum(nil), invertedHandshake.iv())
 }
 
 func ClientHandshake(secret []byte, reader io.Reader) (int, cipher.Stream, cipher.Stream, error) {

mtglib/internal/obfuscated2/handshake_frame.go

@@ -60,3 +60,13 @@ func (h *handshakeFrame) iv() []byte {
 func (h *handshakeFrame) connectionType() []byte {
 	return h.data[handshakeFrameOffsetConnectionType:handshakeFrameOffsetDC]
 }
+
+func (h *handshakeFrame) invert() handshakeFrame {
+	copyFrame := *h
+
+	for i := 0; i < handshakeFrameLenKey+handshakeFrameLenIV; i++ {
+		copyFrame.data[handshakeFrameOffsetKey+i] = h.data[handshakeFrameOffsetConnectionType-1-i]
+	}
+
+	return copyFrame
+}

mtglib/internal/obfuscated2/server_handshake.go

@@ -5,7 +5,7 @@ import (
 	"crypto/rand"
 	"encoding/binary"
 	"fmt"
-	"net"
+	"io"
 )
 
 type serverHandshakeFrame struct {
@@ -13,17 +13,16 @@ type serverHandshakeFrame struct {
 }
 
 func (s *serverHandshakeFrame) decryptor() cipher.Stream {
-	return makeAesCtr(s.key(), s.iv())
+	invertedHandshake := s.invert()
+
+	return makeAesCtr(invertedHandshake.key(), invertedHandshake.iv())
 }
 
 func (s *serverHandshakeFrame) encryptor() cipher.Stream {
-	arr := serverHandshakeFrame{}
+	return makeAesCtr(s.key(), s.iv())
-	invertByteSlices(arr.data[:], s.data[:])
-
-	return makeAesCtr(arr.key(), arr.iv())
 }
 
-func ServerHandshake(conn net.Conn) (cipher.Stream, cipher.Stream, error) {
+func ServerHandshake(writer io.Writer) (cipher.Stream, cipher.Stream, error) {
 	handshake := generateServerHanshakeFrame()
 	copyHandshake := handshake
 	encryptor := handshake.encryptor()
@@ -33,7 +32,7 @@ func ServerHandshake(conn net.Conn) (cipher.Stream, cipher.Stream, error) {
 	copy(handshake.key(), copyHandshake.key())
 	copy(handshake.iv(), copyHandshake.iv())
 
-	if _, err := conn.Write(handshake.data[:]); err != nil {
+	if _, err := writer.Write(handshake.data[:]); err != nil {
 		return nil, nil, fmt.Errorf("cannot send a handshake frame to telegram: %w", err)
 	}
 
@@ -48,16 +47,16 @@ func generateServerHanshakeFrame() serverHandshakeFrame {
 			panic(err)
 		}
 
-		if frame.data[0] == 0xef {
+		if frame.data[0] == 0xef { // nolint: gomnd // taken from tg sources
 			continue
 		}
 
 		switch binary.LittleEndian.Uint32(frame.data[:4]) {
-		case 0x44414548, 0x54534f50, 0x20544547, 0x4954504f, 0xeeeeeeee:
+		case 0x44414548, 0x54534f50, 0x20544547, 0x4954504f, 0xeeeeeeee: // nolint: gomnd // taken from tg sources
 			continue
 		}
 
-		if (frame.data[4] | frame.data[5] | frame.data[6] | frame.data[7]) == 0 {
+		if frame.data[4]|frame.data[5]|frame.data[6]|frame.data[7] == 0 {
 			continue
 		}
 

mtglib/internal/obfuscated2/utils.go

@@ -13,11 +13,3 @@ func makeAesCtr(key, iv []byte) cipher.Stream {
 
 	return cipher.NewCTR(block, iv)
 }
-
-func invertByteSlices(dst, src []byte) {
-	lenDst := len(dst) - 1
-
-	for i, v := range src {
-		dst[lenDst-i] = v
-	}
-}

mtglib/internal/relay/conn.go

@@ -0,0 +1,33 @@
+package relay
+
+import "io"
+
+type conn struct {
+	io.ReadWriteCloser
+
+	relay *Relay
+}
+
+func (c conn) Read(p []byte) (int, error) {
+	ctx := c.relay.ctx
+	n, err := c.ReadWriteCloser.Read(p)
+
+	select {
+	case <-ctx.Done():
+	case c.relay.tickChannel <- struct{}{}:
+	}
+
+	return n, err // nolint: wrapcheck
+}
+
+func (c conn) Write(p []byte) (int, error) {
+	ctx := c.relay.ctx
+	n, err := c.ReadWriteCloser.Write(p)
+
+	select {
+	case <-ctx.Done():
+	case c.relay.tickChannel <- struct{}{}:
+	}
+
+	return n, err // nolint: wrapcheck
+}

mtglib/internal/relay/pools.go

@@ -0,0 +1,46 @@
+package relay
+
+import (
+	"context"
+	"sync"
+	"time"
+)
+
+var relayPool = sync.Pool{
+	New: func() interface{} {
+		return &Relay{
+			tickChannel:  make(chan struct{}),
+			errorChannel: make(chan error, 1),
+		}
+	},
+}
+
+func AcquireRelay(ctx context.Context, logger Logger, bufferSize int, idleTimeout time.Duration) *Relay {
+	ctx, cancel := context.WithCancel(ctx)
+
+	r := relayPool.Get().(*Relay)
+	r.ctx = ctx
+	r.ctxCancel = cancel
+	r.logger = logger
+	r.tickTimeout = idleTimeout
+
+	if len(r.eastBuffer) != bufferSize {
+		r.eastBuffer = make([]byte, bufferSize)
+	}
+
+	if len(r.westBuffer) != bufferSize {
+		r.westBuffer = make([]byte, bufferSize)
+	}
+
+	return r
+}
+
+func ReleaseRelay(r *Relay) {
+	r.ctxCancel()
+
+	r.ctx = nil
+	r.ctxCancel = nil
+	r.logger = nil
+
+	relayPool.Put(r)
+}

mtglib/internal/relay/relay.go

@@ -0,0 +1,100 @@
+package relay
+
+import (
+	"context"
+	"io"
+	"sync"
+	"time"
+)
+
+type Relay struct {
+	ctx          context.Context
+	ctxCancel    context.CancelFunc
+	logger       Logger
+	eastBuffer   []byte
+	westBuffer   []byte
+	tickChannel  chan struct{}
+	errorChannel chan error
+	tickTimeout  time.Duration
+}
+
+func (r *Relay) Process(eastConn, westConn io.ReadWriteCloser) error {
+	eastConn = conn{
+		ReadWriteCloser: eastConn,
+		relay:           r,
+	}
+	westConn = conn{
+		ReadWriteCloser: westConn,
+		relay:           r,
+	}
+
+	defer func() {
+		r.ctxCancel()
+		eastConn.Close()
+		westConn.Close()
+	}()
+
+	go r.runObserver()
+
+	wg := &sync.WaitGroup{}
+	wg.Add(2) // nolint: gomnd
+
+	go r.transmit(eastConn, westConn, r.westBuffer, "west", wg)
+
+	r.transmit(westConn, eastConn, r.eastBuffer, "east", wg)
+
+	wg.Wait()
+
+	return <-r.errorChannel
+}
+
+func (r *Relay) transmit(src io.ReadCloser, dst io.WriteCloser,
+	buffer []byte, direction string, wg *sync.WaitGroup) {
+	defer func() {
+		wg.Done()
+		src.Close()
+		dst.Close()
+	}()
+
+	if _, err := io.CopyBuffer(dst, src, buffer); err != nil {
+		r.logger.Printf("error '%v' happened on direction %s", err, direction)
+
+		select {
+		case <-r.ctx.Done():
+		case r.errorChannel <- err:
+		default:
+		}
+	}
+}
+
+func (r *Relay) runObserver() {
+	ticker := time.NewTicker(time.Second)
+
+	defer func() {
+		ticker.Stop()
+
+		select {
+		case <-ticker.C:
+		default:
+		}
+	}()
+
+	lastTickAt := time.Now()
+	ctx := r.ctx
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-r.tickChannel:
+			lastTickAt = time.Now()
+		case <-ticker.C:
+			if time.Since(lastTickAt) > r.tickTimeout {
+				r.logger.Printf("exit due to a timeout")
+				r.ctxCancel()
+
+				return
+			}
+		}
+	}
+}

mtglib/proxy.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/9seconds/mtg/v2/mtglib/internal/obfuscated2"
+	"github.com/9seconds/mtg/v2/mtglib/internal/relay"
 	"github.com/9seconds/mtg/v2/mtglib/internal/telegram"
 	"github.com/panjf2000/ants/v2"
 )
@@ -19,6 +20,7 @@ type Proxy struct {
 	streamWaitGroup sync.WaitGroup
 
 	idleTimeout time.Duration
+	bufferSize  int
 	workerPool  *ants.PoolWithFunc
 	telegram    *telegram.Telegram
 
@@ -64,6 +66,13 @@ func (p *Proxy) ServeConn(conn net.Conn) {
 
 		return
 	}
+
+	rel := relay.AcquireRelay(ctx, p.logger.Named("relay"), p.bufferSize, p.idleTimeout)
+	defer relay.ReleaseRelay(rel)
+
+	if err := rel.Process(ctx.clientConn, ctx.telegramConn); err != nil {
+		p.logger.DebugError("relay has been finished", err)
+	}
 }
 
 func (p *Proxy) Serve(listener net.Listener) error {
@@ -185,6 +194,11 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) { // nolint: cyclop
 		idleTimeout = DefaultIdleTimeout
 	}
 
+	bufferSize := opts.BufferSize
+	if bufferSize < 1 {
+		bufferSize = DefaultBufferSize
+	}
+
 	ctx, cancel := context.WithCancel(context.Background())
 	proxy := &Proxy{
 		ctx:             ctx,
@@ -195,6 +209,7 @@ func NewProxy(opts ProxyOpts) (*Proxy, error) { // nolint: cyclop
 		eventStream:     opts.EventStream,
 		logger:          opts.Logger.Named("proxy"),
 		idleTimeout:     idleTimeout,
+		bufferSize:      int(bufferSize),
 		telegram:        tg,
 	}