Propagate to middle client

client/middle.go

@@ -0,0 +1,25 @@
+package client
+
+import (
+	"net"
+
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
+	mtwrappers "github.com/9seconds/mtg/mtproto/wrappers"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+func MiddleInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, wrappers.ReadWriteCloserWithAddr, error) {
+	opts, newConn, err := DirectInit(conn, conf)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if opts.ConnectionType == mtproto.ConnectionTypeAbridged {
+		newConn = mtwrappers.NewAbridgedRWC(newConn, opts)
+	} else {
+		newConn = mtwrappers.NewIntermediateRWC(newConn, opts)
+	}
+
+	return opts, newConn, nil
+}

main.go

@@ -4,7 +4,6 @@ package main
 
 import (
 	"encoding/json"
-	"fmt"
 	"io"
 	"math/rand"
 	"os"
@@ -16,8 +15,7 @@ import (
 	kingpin "gopkg.in/alecthomas/kingpin.v2"
 
 	"github.com/9seconds/mtg/config"
-	"github.com/9seconds/mtg/mtproto"
-	"github.com/9seconds/mtg/telegram"
+	"github.com/9seconds/mtg/proxy"
 	"github.com/juju/errors"
 )
 
@@ -115,29 +113,15 @@ func main() {
 		atom,
 	)).Sugar()
 
-	tg := telegram.NewMiddleTelegram(conf, logger)
-	connOpts := &mtproto.ConnectionOpts{
-		DC:              int16(1),
-		ConnectionType:  mtproto.ConnectionTypeAbridged,
-		ConnectionProto: mtproto.ConnectionProtocolIPv4,
+	stat := proxy.NewStats(conf)
+	go stat.Serve()
+
+	srv := proxy.NewServer(conf, logger, stat)
+	printURLs(conf.GetURLs())
+
+	if err := srv.Serve(); err != nil {
+		logger.Fatal(err.Error())
 	}
-
-	sock, err := tg.Dial(connOpts)
-	if err != nil {
-		panic(err)
-	}
-	_, err = tg.Init(connOpts, sock)
-	fmt.Println(err)
-
-	// stat := proxy.NewStats(conf)
-	// go stat.Serve()
-
-	// srv := proxy.NewServer(conf, logger, stat)
-	// printURLs(conf.GetURLs())
-
-	// if err := srv.Serve(); err != nil {
-	// 	logger.Fatal(err.Error())
-	// }
 }
 
 func setRLimit() (err error) {

proxy/server.go

@@ -139,11 +139,19 @@ func (s *Server) pipe(dst io.Writer, src io.Reader, wait *sync.WaitGroup) {
 
 // NewServer creates new instance of MTPROTO proxy.
 func NewServer(conf *config.Config, logger *zap.SugaredLogger, stat *Stats) *Server {
+	clientInit := client.DirectInit
+	tg := telegram.NewDirectTelegram
+
+	if len(conf.AdTag) > 0 {
+		clientInit = client.MiddleInit
+		tg = telegram.NewMiddleTelegram
+	}
+
 	return &Server{
 		conf:       conf,
 		logger:     logger,
 		stats:      stat,
-		tg:         telegram.NewDirectTelegram(conf),
-		clientInit: client.DirectInit,
+		tg:         tg(conf, logger),
+		clientInit: clientInit,
 	}
 }

telegram/direct.go

@@ -4,6 +4,7 @@ import (
 	"net"
 
 	"github.com/juju/errors"
+	"go.uber.org/zap"
 
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
@@ -55,7 +56,7 @@ func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts, conn wrappers.Re
 
 // NewDirectTelegram returns Telegram instance which connects directly
 // to Telegram bypassing middleproxies.
-func NewDirectTelegram(conf *config.Config) Telegram {
+func NewDirectTelegram(conf *config.Config, _ *zap.SugaredLogger) Telegram {
 	return &directTelegram{baseTelegram{
 		dialer: tgDialer{
 			Dialer: net.Dialer{Timeout: telegramDialTimeout},