From a3362c7ea4c7b8eec6b37dfb542c81c310ddae18 Mon Sep 17 00:00:00 2001 From: 9seconds Date: Tue, 23 Mar 2021 15:14:36 +0300 Subject: [PATCH] Proxy is working in a simple mode now --- mtglib/conns.go | 4 +- .../internal/obfuscated2/client_handshake.go | 9 +- .../internal/obfuscated2/handshake_frame.go | 10 ++ .../internal/obfuscated2/server_handshake.go | 21 ++-- mtglib/internal/obfuscated2/utils.go | 8 -- mtglib/internal/relay/conn.go | 33 ++++++ mtglib/internal/relay/pools.go | 46 ++++++++ mtglib/internal/relay/relay.go | 100 ++++++++++++++++++ mtglib/proxy.go | 15 +++ 9 files changed, 220 insertions(+), 26 deletions(-) create mode 100644 mtglib/internal/relay/conn.go create mode 100644 mtglib/internal/relay/pools.go create mode 100644 mtglib/internal/relay/relay.go diff --git a/mtglib/conns.go b/mtglib/conns.go index 985cad0..e30cf0d 100644 --- a/mtglib/conns.go +++ b/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 } diff --git a/mtglib/internal/obfuscated2/client_handshake.go b/mtglib/internal/obfuscated2/client_handshake.go index 8601766..89cc396 100644 --- a/mtglib/internal/obfuscated2/client_handshake.go +++ b/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{} - invertByteSlices(arr.data[:], c.data[:]) + invertedHandshake := c.invert() hasher := acquireSha256Hasher() defer releaseSha256Hasher(hasher) - hasher.Write(arr.key()) // nolint: errcheck - hasher.Write(secret) // nolint: errcheck + hasher.Write(invertedHandshake.key()) // 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) { diff --git a/mtglib/internal/obfuscated2/handshake_frame.go b/mtglib/internal/obfuscated2/handshake_frame.go index 1309981..82327f1 100644 --- a/mtglib/internal/obfuscated2/handshake_frame.go +++ b/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 +} diff --git a/mtglib/internal/obfuscated2/server_handshake.go b/mtglib/internal/obfuscated2/server_handshake.go index 3b69d6a..9712fde 100644 --- a/mtglib/internal/obfuscated2/server_handshake.go +++ b/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{} - invertByteSlices(arr.data[:], s.data[:]) - - return makeAesCtr(arr.key(), arr.iv()) + return makeAesCtr(s.key(), s.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 } diff --git a/mtglib/internal/obfuscated2/utils.go b/mtglib/internal/obfuscated2/utils.go index f32dde0..a3fdcf9 100644 --- a/mtglib/internal/obfuscated2/utils.go +++ b/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 - } -} diff --git a/mtglib/internal/relay/conn.go b/mtglib/internal/relay/conn.go new file mode 100644 index 0000000..211a32a --- /dev/null +++ b/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 +} diff --git a/mtglib/internal/relay/pools.go b/mtglib/internal/relay/pools.go new file mode 100644 index 0000000..dc51c0f --- /dev/null +++ b/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) +} diff --git a/mtglib/internal/relay/relay.go b/mtglib/internal/relay/relay.go new file mode 100644 index 0000000..2a6e1b5 --- /dev/null +++ b/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 + } + } + } +} diff --git a/mtglib/proxy.go b/mtglib/proxy.go index 6678513..43b1bbe 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/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, }