Rework for simplier proxy

2026-08-31 16:04:02 +03:00 · 2018-07-08 08:04:29 +03:00
parent cd63483503
commit a0ba89b105
12 changed files with 154 additions and 195 deletions
@@ -1,64 +0,0 @@
-package proxy
-
-import (
-	"context"
-	"io"
-	"net"
-	"sync"
-
-	"github.com/juju/errors"
-
-	"github.com/9seconds/mtg/client"
-	"github.com/9seconds/mtg/config"
-	"github.com/9seconds/mtg/mtproto"
-	"github.com/9seconds/mtg/telegram"
-	"github.com/9seconds/mtg/wrappers"
-)
-
-func NewProxyDirect(conf *config.Config) *Proxy {
-	tg := telegram.NewDirectTelegram(conf)
-
-	return &Proxy{
-		conf: conf,
-		acceptCallback: func(ctx context.Context, cancel context.CancelFunc, clientSocket net.Conn,
-			connID string, wait *sync.WaitGroup, conf *config.Config) (io.Closer, io.Closer, error) {
-			client, opts, err := client.DirectInit(ctx, cancel, clientSocket, connID, conf)
-			if err != nil {
-				return nil, nil, errors.Annotate(err, "Cannot initialize client connection")
-			}
-
-			server, err := directTelegramStream(ctx, cancel, opts, connID, tg)
-			if err != nil {
-				return client, nil, errors.Annotate(err, "Cannot initialize telegram connection")
-			}
-
-			wait.Add(2)
-
-			go directPipe(client, server, wait)
-			go directPipe(server, client, wait)
-
-			return client, server, nil
-		},
-	}
-}
-
-func directTelegramStream(ctx context.Context, cancel context.CancelFunc, opts *mtproto.ConnectionOpts,
-	connID string, tg *telegram.DirectTelegram) (wrappers.WrapStreamReadWriteCloser, error) {
-	streamConn, err := tg.Dial(connID, opts)
-	if err != nil {
-		return nil, errors.Annotate(err, "Cannot dial to Telegram")
-	}
-	streamConn = wrappers.NewCtx(ctx, cancel, streamConn)
-
-	packetConn, err := tg.Init(opts, streamConn)
-	if err != nil {
-		return nil, errors.Annotate(err, "Cannot handshake telegram")
-	}
-
-	return packetConn, nil
-}
-
-func directPipe(src io.Reader, dst io.Writer, wait *sync.WaitGroup) {
-	defer wait.Done()
-	io.Copy(dst, src)
-}
@@ -1,76 +0,0 @@
-package proxy
-
-import (
-	"context"
-	"io"
-	"net"
-	"sync"
-
-	"github.com/juju/errors"
-
-	"github.com/9seconds/mtg/client"
-	"github.com/9seconds/mtg/config"
-	"github.com/9seconds/mtg/mtproto"
-	"github.com/9seconds/mtg/telegram"
-	"github.com/9seconds/mtg/wrappers"
-)
-
-func NewProxyMiddle(conf *config.Config) *Proxy {
-	tg := telegram.NewMiddleTelegram(conf)
-
-	return &Proxy{
-		conf: conf,
-		acceptCallback: func(ctx context.Context, cancel context.CancelFunc, clientSocket net.Conn,
-			connID string, wait *sync.WaitGroup, conf *config.Config) (io.Closer, io.Closer, error) {
-			client, opts, err := client.MiddleInit(ctx, cancel, clientSocket, connID, conf)
-			if err != nil {
-				return nil, nil, errors.Annotate(err, "Cannot initialize client connection")
-			}
-
-			server, err := middleTelegramStream(ctx, cancel, opts, connID, tg)
-			if err != nil {
-				return client, nil, errors.Annotate(err, "Cannot initialize telegram connection")
-			}
-
-			wait.Add(2)
-
-			go middlePipe(client, server, wait, &opts.ReadHacks)
-			go middlePipe(server, client, wait, &opts.WriteHacks)
-
-			return client, server, nil
-		},
-	}
-}
-
-func middleTelegramStream(ctx context.Context, cancel context.CancelFunc, opts *mtproto.ConnectionOpts,
-	connID string, tg *telegram.MiddleTelegram) (wrappers.WrapPacketReadWriteCloser, error) {
-	streamConn, err := tg.Dial(connID, opts)
-	if err != nil {
-		return nil, errors.Annotate(err, "Cannot dial to Telegram")
-	}
-	streamConn = wrappers.NewCtx(ctx, cancel, streamConn)
-
-	packetConn, err := tg.Init(opts, streamConn)
-	if err != nil {
-		return nil, errors.Annotate(err, "Cannot handshake telegram")
-	}
-
-	return packetConn, nil
-}
-
-func middlePipe(src wrappers.WrapPacketReader, dst wrappers.WrapPacketWriter, wait *sync.WaitGroup, hacks *mtproto.Hacks) {
-	defer wait.Done()
-
-	for {
-		hacks.SimpleAck = false
-		hacks.QuickAck = false
-
-		packet, err := src.Read()
-		if err != nil {
-			return
-		}
-		if _, err = dst.Write(packet); err != nil {
-			return
-		}
-	}
-}
@@ -10,14 +10,17 @@ import (
 	uuid "github.com/satori/go.uuid"
 	"go.uber.org/zap"

+	"github.com/9seconds/mtg/client"
 	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/telegram"
+	"github.com/9seconds/mtg/wrappers"
 )

-type proxyAcceptCallback func(context.Context, context.CancelFunc, net.Conn, string, *sync.WaitGroup, *config.Config) (io.Closer, io.Closer, error)
-
 type Proxy struct {
-	conf           *config.Config
-	acceptCallback proxyAcceptCallback
+	clientInit client.Init
+	tg         telegram.Telegram
+	conf       *config.Config
 }

 func (p *Proxy) Serve() error {
@@ -50,20 +53,33 @@ func (p *Proxy) accept(conn net.Conn) {
 	log.Infow("Client connected", "addr", conn.RemoteAddr())

 	ctx, cancel := context.WithCancel(context.Background())
-	wait := &sync.WaitGroup{}
-
-	client, server, err := p.acceptCallback(ctx, cancel, conn, connID, wait, p.conf)
-	defer func() {
-		if client != nil {
-			client.Close()
-		}
-		if server != nil {
-			server.Close()
-		}
-	}()
+	client, opts, err := p.clientInit(ctx, cancel, conn, connID, p.conf)
 	if err != nil {
-		log.Errorw("Cannot initialize connection", "error", err)
-		cancel()
+		log.Errorw("Cannot initialize client connection", "error", err)
+		return
+	}
+	defer client.(wrappers.WrapCloser).Close()
+
+	server, err := p.getTelegramConn(ctx, cancel, opts, connID)
+	if err != nil {
+		log.Errorw("Cannot initialize server connection", "error", err)
+		return
+	}
+	defer server.(wrappers.WrapCloser).Close()
+
+	wait := &sync.WaitGroup{}
+	wait.Add(2)
+
+	if p.conf.UseMiddleProxy() {
+		clientPacket := client.(wrappers.WrapPacketReadWriteCloser)
+		serverPacket := server.(wrappers.WrapPacketReadWriteCloser)
+		go p.middlePipe(clientPacket, serverPacket, wait, &opts.ReadHacks)
+		go p.middlePipe(serverPacket, clientPacket, wait, &opts.WriteHacks)
+	} else {
+		clientStream := client.(wrappers.WrapStreamReadWriteCloser)
+		serverStream := server.(wrappers.WrapStreamReadWriteCloser)
+		go p.directPipe(clientStream, serverStream, wait)
+		go p.directPipe(serverStream, clientStream, wait)
 	}

 	<-ctx.Done()
@@ -71,3 +87,61 @@ func (p *Proxy) accept(conn net.Conn) {

 	log.Infow("Client disconnected", "addr", conn.RemoteAddr())
 }
+
+func (p *Proxy) getTelegramConn(ctx context.Context, cancel context.CancelFunc, opts *mtproto.ConnectionOpts,
+	connID string) (wrappers.Wrap, error) {
+	streamConn, err := p.tg.Dial(connID, opts)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot dial to Telegram")
+	}
+	streamConn = wrappers.NewCtx(ctx, cancel, streamConn)
+
+	packetConn, err := p.tg.Init(opts, streamConn)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot handshake telegram")
+	}
+
+	return packetConn, nil
+}
+
+func (p *Proxy) middlePipe(src wrappers.WrapPacketReader, dst wrappers.WrapPacketWriter, wait *sync.WaitGroup, hacks *mtproto.Hacks) {
+	defer wait.Done()
+
+	for {
+		hacks.SimpleAck = false
+		hacks.QuickAck = false
+
+		packet, err := src.Read()
+		if err != nil {
+			return
+		}
+		if _, err = dst.Write(packet); err != nil {
+			return
+		}
+	}
+}
+
+func (p *Proxy) directPipe(src io.Reader, dst io.Writer, wait *sync.WaitGroup) {
+	defer wait.Done()
+	io.Copy(dst, src)
+
+}
+
+func NewProxy(conf *config.Config) *Proxy {
+	var clientInit client.Init
+	var tg telegram.Telegram
+
+	if conf.UseMiddleProxy() {
+		clientInit = client.MiddleInit
+		tg = telegram.NewMiddleTelegram(conf)
+	} else {
+		clientInit = client.DirectInit
+		tg = telegram.NewDirectTelegram(conf)
+	}
+
+	return &Proxy{
+		conf:       conf,
+		clientInit: clientInit,
+		tg:         tg,
+	}
+}