Rework for simplier proxy

client/client.go

@@ -0,0 +1,12 @@
+package client
+
+import (
+	"context"
+	"net"
+
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+type Init func(context.Context, context.CancelFunc, net.Conn, string, *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error)

client/direct.go

@@ -16,7 +16,7 @@ import (
 const handshakeTimeout = 10 * time.Second
 
 func DirectInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn, connID string,
-	conf *config.Config) (wrappers.WrapStreamReadWriteCloser, *mtproto.ConnectionOpts, error) {
+	conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) {
 	if err := config.SetSocketOptions(socket); err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot set socket options")
 	}

client/middle.go

@@ -10,15 +10,16 @@ import (
 )
 
 func MiddleInit(ctx context.Context, cancel context.CancelFunc, socket net.Conn, connID string,
-	conf *config.Config) (wrappers.WrapPacketReadWriteCloser, *mtproto.ConnectionOpts, error) {
+	conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) {
 	conn, opts, err := DirectInit(ctx, cancel, socket, connID, conf)
 	if err != nil {
 		return nil, nil, err
 	}
+	connStream := conn.(wrappers.WrapStreamReadWriteCloser)
 
-	newConn := wrappers.NewMTProtoAbridged(conn, opts)
+	newConn := wrappers.NewMTProtoAbridged(connStream, opts)
 	if opts.ConnectionType != mtproto.ConnectionTypeAbridged {
-		newConn = wrappers.NewMTProtoIntermediate(conn, opts)
+		newConn = wrappers.NewMTProtoIntermediate(connStream, opts)
 	}
 
 	opts.ConnectionProto = mtproto.ConnectionProtocolIPv4

config/config.go

@@ -58,6 +58,10 @@ func (c *Config) StatAddr() string {
 	return getAddr(c.StatsIP, c.StatsPort)
 }
 
+func (c *Config) UseMiddleProxy() bool {
+	return len(c.AdTag) > 0
+}
+
 // GetURLs returns configured IPURLs instance with links to this server.
 func (c *Config) GetURLs() IPURLs {
 	urls := IPURLs{}

main.go

@@ -115,17 +115,15 @@ func main() {
 	zap.ReplaceGlobals(logger)
 	defer logger.Sync()
 
-	var server *proxy.Proxy
+	if conf.UseMiddleProxy() {
-	if len(conf.AdTag) == 0 {
-		zap.S().Infow("Use direct connection to Telegram")
-		server = proxy.NewProxyDirect(conf)
-	} else {
 		zap.S().Infow("Use middle proxy connection to Telegram")
-		server = proxy.NewProxyMiddle(conf)
+	} else {
+		zap.S().Infow("Use direct connection to Telegram")
 	}
 
 	printURLs(conf.GetURLs())
 
+	server := proxy.NewProxy(conf)
 	if err := server.Serve(); err != nil {
 		zap.S().Fatalw("Server stopped", "error", err)
 	}

proxy/direct.go

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

proxy/middle.go

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

proxy/proxy.go

@@ -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
+	clientInit client.Init
-	acceptCallback proxyAcceptCallback
+	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, opts, err := p.clientInit(ctx, cancel, conn, connID, p.conf)
-
-	client, server, err := p.acceptCallback(ctx, cancel, conn, connID, wait, p.conf)
-	defer func() {
-		if client != nil {
-			client.Close()
-		}
-		if server != nil {
-			server.Close()
-		}
-	}()
 	if err != nil {
-		log.Errorw("Cannot initialize connection", "error", err)
+		log.Errorw("Cannot initialize client connection", "error", err)
-		cancel()
+		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,
+	}
+}

telegram/direct.go

@@ -43,7 +43,7 @@ func (t *DirectTelegram) Dial(connID string, connOpts *mtproto.ConnectionOpts) (
 	return t.baseTelegram.dial(dc-1, connID, connOpts.ConnectionProto)
 }
 
-func (t *DirectTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.WrapStreamReadWriteCloser) (wrappers.WrapStreamReadWriteCloser, error) {
+func (t *DirectTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.WrapStreamReadWriteCloser) (wrappers.Wrap, error) {
 	obfs2, frame := obfuscated2.MakeTelegramObfuscated2Frame(connOpts)
 
 	if _, err := conn.Write(frame); err != nil {
@@ -55,7 +55,7 @@ func (t *DirectTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.Wr
 
 // NewDirectTelegram returns Telegram instance which connects directly
 // to Telegram bypassing middleproxies.
-func NewDirectTelegram(conf *config.Config) *DirectTelegram {
+func NewDirectTelegram(conf *config.Config) Telegram {
 	return &DirectTelegram{baseTelegram{
 		dialer: tgDialer{
 			Dialer: net.Dialer{Timeout: telegramDialTimeout},

telegram/middle.go

@@ -19,32 +19,7 @@ type MiddleTelegram struct {
 	conf *config.Config
 }
 
-func NewMiddleTelegram(conf *config.Config) *MiddleTelegram {
+func (t *MiddleTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.WrapStreamReadWriteCloser) (wrappers.Wrap, error) {
-	tg := &MiddleTelegram{
-		middleTelegramCaller: middleTelegramCaller{
-			baseTelegram: baseTelegram{
-				dialer: tgDialer{
-					Dialer: net.Dialer{Timeout: telegramDialTimeout},
-					conf:   conf,
-				},
-			},
-			httpClient: &http.Client{
-				Timeout: middleTelegramHTTPClientTimeout,
-			},
-			dialerMutex: &sync.RWMutex{},
-		},
-		conf: conf,
-	}
-
-	if err := tg.update(); err != nil {
-		panic(err)
-	}
-	go tg.autoUpdate()
-
-	return tg
-}
-
-func (t *MiddleTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.WrapStreamReadWriteCloser) (wrappers.WrapPacketReadWriteCloser, error) {
 	rpcNonceConn := wrappers.NewMTProtoFrame(conn, rpc.SeqNoNonce)
 
 	rpcNonceReq, err := t.sendRPCNonceRequest(rpcNonceConn)
@@ -125,3 +100,28 @@ func (t *MiddleTelegram) receiveRPCHandshakeResponse(conn wrappers.WrapPacketRea
 
 	return rpcHandshakeResp, nil
 }
+
+func NewMiddleTelegram(conf *config.Config) Telegram {
+	tg := &MiddleTelegram{
+		middleTelegramCaller: middleTelegramCaller{
+			baseTelegram: baseTelegram{
+				dialer: tgDialer{
+					Dialer: net.Dialer{Timeout: telegramDialTimeout},
+					conf:   conf,
+				},
+			},
+			httpClient: &http.Client{
+				Timeout: middleTelegramHTTPClientTimeout,
+			},
+			dialerMutex: &sync.RWMutex{},
+		},
+		conf: conf,
+	}
+
+	if err := tg.update(); err != nil {
+		panic(err)
+	}
+	go tg.autoUpdate()
+
+	return tg
+}

telegram/telegram.go

@@ -9,6 +9,11 @@ import (
 	"github.com/9seconds/mtg/wrappers"
 )
 
+type Telegram interface {
+	Dial(string, *mtproto.ConnectionOpts) (wrappers.WrapStreamReadWriteCloser, error)
+	Init(*mtproto.ConnectionOpts, wrappers.WrapStreamReadWriteCloser) (wrappers.Wrap, error)
+}
+
 type baseTelegram struct {
 	dialer tgDialer
 

wrappers/wrap.go

@@ -20,6 +20,11 @@ type WrapWriter interface {
 	Wrap
 }
 
+type WrapCloser interface {
+	io.Closer
+	Wrap
+}
+
 type WrapWriteCloser interface {
 	io.Closer
 	WrapWriter