Merge pull request #10 from 9seconds/channels

Promoted channels

Gopkg.toml

@@ -44,3 +44,7 @@
 [[constraint]]
   name = "github.com/satori/go.uuid"
   version = "1.2.0"
+
+[[constraint]]
+  branch = "master"
+  name = "github.com/dustin/go-humanize"

README.md

@@ -38,6 +38,18 @@ mtg is an implementation in golang which is intended to be:
 * **No management WebUI**
   This is an implementation of simple lightweight proxy. I won't do that.
 
+This proxy supports 2 modes of work: direct connection to Telegram and
+promoted channel mode. If you do not need promoted channels, I would
+recommend you to go with direct mode: this is way more robust.
+
+To run proxy in direct mode, all you need to do is just provide a
+secret. If you do not provide ADTag as a second parameter, promoted
+channels mode won't be activated.
+
+To get promoted channel, please contact
+[@MTProxybot|https://t.me/MTProxybot] and provide generated adtag as a
+second parameter.
+
 
 # How to build
 

client/client.go

@@ -1,12 +1,12 @@
 package client
 
 import (
-	"io"
 	"net"
 
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/wrappers"
 )
 
-// Init has to initialize client connection based on given config.
-type Init func(net.Conn, *config.Config) (*mtproto.ConnectionOpts, io.ReadWriteCloser, error)
+// Init defines common method for initializing client connections.
+type Init func(net.Conn, string, *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error)

client/direct.go

@@ -1,7 +1,6 @@
 package client
 
 import (
-	"io"
 	"net"
 	"time"
 
@@ -15,28 +14,42 @@ import (
 
 const (
 	handshakeTimeout = 10 * time.Second
+	readBufferSize   = 64 * 1024
+	writeBufferSize  = 64 * 1024
 )
 
-// DirectInit initializes client to access Telegram bypassing middleproxies.
-func DirectInit(conn net.Conn, conf *config.Config) (*mtproto.ConnectionOpts, io.ReadWriteCloser, error) {
-	if err := config.SetSocketOptions(conn); err != nil {
-		return nil, nil, errors.Annotate(err, "Cannot set socket options")
+// DirectInit initializes client connection for proxy which connects to
+// Telegram directly.
+func DirectInit(socket net.Conn, connID string, conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) {
+	tcpSocket := socket.(*net.TCPConn)
+	if err := tcpSocket.SetNoDelay(false); err != nil {
+		return nil, nil, errors.Annotate(err, "Cannot disable NO_DELAY to client socket")
+	}
+	if err := tcpSocket.SetReadBuffer(readBufferSize); err != nil {
+		return nil, nil, errors.Annotate(err, "Cannot set read buffer size of client socket")
+	}
+	if err := tcpSocket.SetWriteBuffer(writeBufferSize); err != nil {
+		return nil, nil, errors.Annotate(err, "Cannot set write buffer size of client socket")
 	}
 
-	conn.SetReadDeadline(time.Now().Add(handshakeTimeout)) // nolint: errcheck
-	frame, err := obfuscated2.ExtractFrame(conn)
-	conn.SetReadDeadline(time.Time{}) // nolint: errcheck
+	socket.SetReadDeadline(time.Now().Add(handshakeTimeout)) // nolint: errcheck
+	frame, err := obfuscated2.ExtractFrame(socket)
 	if err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot extract frame")
 	}
-	defer obfuscated2.ReturnFrame(frame)
+	socket.SetReadDeadline(time.Time{}) // nolint: errcheck
 
+	conn := wrappers.NewConn(socket, connID, wrappers.ConnPurposeClient, conf.PublicIPv4, conf.PublicIPv6)
 	obfs2, connOpts, err := obfuscated2.ParseObfuscated2ClientFrame(conf.Secret, frame)
 	if err != nil {
 		return nil, nil, errors.Annotate(err, "Cannot parse obfuscated frame")
 	}
+	connOpts.ConnectionProto = mtproto.ConnectionProtocolAny
+	connOpts.ClientAddr = conn.RemoteAddr()
+
+	conn = wrappers.NewStreamCipher(conn, obfs2.Encryptor, obfs2.Decryptor)
 
-	socket := wrappers.NewStreamCipherRWC(conn, obfs2.Encryptor, obfs2.Decryptor)
+	conn.Logger().Infow("Client connection initialized")
 
-	return connOpts, socket, nil
+	return conn, connOpts, nil
 }

client/middle.go

@@ -0,0 +1,31 @@
+package client
+
+import (
+	"net"
+
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+// MiddleInit initializes client connection for proxy which has to
+// support promoted channels, connect to Telegram middle proxies etc.
+func MiddleInit(socket net.Conn, connID string, conf *config.Config) (wrappers.Wrap, *mtproto.ConnectionOpts, error) {
+	conn, opts, err := DirectInit(socket, connID, conf)
+	if err != nil {
+		return nil, nil, err
+	}
+	connStream := conn.(wrappers.StreamReadWriteCloser)
+
+	newConn := wrappers.NewMTProtoAbridged(connStream, opts)
+	if opts.ConnectionType != mtproto.ConnectionTypeAbridged {
+		newConn = wrappers.NewMTProtoIntermediate(connStream, opts)
+	}
+
+	opts.ConnectionProto = mtproto.ConnectionProtocolIPv4
+	if socket.LocalAddr().(*net.TCPAddr).IP.To4() == nil {
+		opts.ConnectionProto = mtproto.ConnectionProtocolIPv6
+	}
+
+	return newConn, opts, err
+}

config/config.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"net"
 	"strconv"
-	"time"
 
 	"github.com/juju/errors"
 )
@@ -14,9 +13,6 @@ import (
 const (
 	BufferWriteSize = 32 * 1024
 	BufferReadSize  = 32 * 1024
-	BufferSizeCopy  = 32 * 1024
-
-	keepAlivePeriod = 20 * time.Second
 )
 
 // Config represents common configuration of mtg.
@@ -35,6 +31,7 @@ type Config struct {
 	StatsIP    net.IP
 
 	Secret []byte
+	AdTag  []byte
 }
 
 // URLs contains links to the proxy (tg://, t.me) and their QR codes.
@@ -56,27 +53,28 @@ func (c *Config) BindAddr() string {
 	return getAddr(c.BindIP, c.BindPort)
 }
 
-// IPv4Addr returns connection string to ipv6 for mtproto proxy.
-func (c *Config) IPv4Addr() string {
-	return getAddr(c.PublicIPv4, c.PublicIPv4Port)
-}
-
-// IPv6Addr returns connection string to ipv6 for mtproto proxy.
-func (c *Config) IPv6Addr() string {
-	return getAddr(c.PublicIPv6, c.PublicIPv6Port)
-}
-
 // StatAddr returns connection string to the stats API.
 func (c *Config) StatAddr() string {
 	return getAddr(c.StatsIP, c.StatsPort)
 }
 
+// UseMiddleProxy defines if this proxy has to connect middle proxies
+// which supports promoted channels or directly access Telegram.
+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 {
-	return IPURLs{
-		IPv4: getURLs(c.PublicIPv4, c.PublicIPv4Port, c.Secret),
-		IPv6: getURLs(c.PublicIPv6, c.PublicIPv6Port, c.Secret),
+	urls := IPURLs{}
+	if c.PublicIPv4 != nil {
+		urls.IPv4 = getURLs(c.PublicIPv4, c.PublicIPv4Port, c.Secret)
+	}
+	if c.PublicIPv6 != nil {
+		urls.IPv6 = getURLs(c.PublicIPv6, c.PublicIPv6Port, c.Secret)
 	}
+
+	return urls
 }
 
 func getAddr(host fmt.Stringer, port uint16) string {
@@ -91,7 +89,7 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 	publicIPv4 net.IP, PublicIPv4Port uint16,
 	publicIPv6 net.IP, publicIPv6Port uint16,
 	statsIP net.IP, statsPort uint16,
-	secret string) (*Config, error) {
+	secret, adtag string) (*Config, error) {
 	if len(secret) != 32 {
 		return nil, errors.New("Telegram demands secret of length 32")
 	}
@@ -100,28 +98,34 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 		return nil, errors.Annotate(err, "Cannot create config")
 	}
 
+	var adTagBytes []byte
+	if len(adtag) != 0 {
+		adTagBytes, err = hex.DecodeString(adtag)
+		if err != nil {
+			return nil, errors.Annotate(err, "Cannot create config")
+		}
+	}
+
 	if publicIPv4 == nil {
 		publicIPv4, err = getGlobalIPv4()
 		if err != nil {
-			return nil, errors.Errorf("Cannot get public IP")
+			publicIPv4 = nil
+		} else if publicIPv4.To4() == nil {
+			return nil, errors.Errorf("IP %s is not IPv4", publicIPv4.String())
 		}
 	}
-	if publicIPv4.To4() == nil {
-		return nil, errors.Errorf("IP %s is not IPv4", publicIPv4.String())
-	}
 	if PublicIPv4Port == 0 {
 		PublicIPv4Port = bindPort
 	}
 
 	if publicIPv6 == nil {
 		publicIPv6, err = getGlobalIPv6()
 		if err != nil {
-			publicIPv6 = publicIPv4
+			publicIPv6 = nil
+		} else if publicIPv6.To4() != nil {
+			return nil, errors.Errorf("IP %s is not IPv6", publicIPv6.String())
 		}
 	}
-	if publicIPv6.To16() == nil {
-		return nil, errors.Errorf("IP %s is not IPv6", publicIPv6.String())
-	}
 	if publicIPv6Port == 0 {
 		publicIPv6Port = bindPort
 	}
@@ -142,30 +146,8 @@ func NewConfig(debug, verbose bool, // nolint: gocyclo
 		StatsIP:        statsIP,
 		StatsPort:      statsPort,
 		Secret:         secretBytes,
+		AdTag:          adTagBytes,
 	}
 
 	return conf, nil
 }
-
-// SetSocketOptions makes socket keepalive, sets buffer sizes
-func SetSocketOptions(conn net.Conn) error {
-	socket := conn.(*net.TCPConn)
-
-	if err := socket.SetReadBuffer(BufferReadSize); err != nil {
-		return errors.Annotate(err, "Cannot set read buffer size")
-	}
-	if err := socket.SetWriteBuffer(BufferWriteSize); err != nil {
-		return errors.Annotate(err, "Cannot set write buffer size")
-	}
-	if err := socket.SetKeepAlive(true); err != nil {
-		return errors.Annotate(err, "Cannot make socket keepalive")
-	}
-	if err := socket.SetKeepAlivePeriod(keepAlivePeriod); err != nil {
-		return errors.Annotate(err, "Cannot set keepalive period")
-	}
-	if err := socket.SetNoDelay(true); err != nil {
-		return errors.Annotate(err, "Cannot activate nodelay for the socket")
-	}
-
-	return nil
-}

main.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/proxy"
+	"github.com/9seconds/mtg/stats"
 	"github.com/juju/errors"
 )
 
@@ -70,6 +71,7 @@ var (
 			Uint16()
 
 	secret = app.Arg("secret", "Secret of this proxy.").Required().String()
+	adtag  = app.Arg("adtag", "ADTag of the proxy.").String()
 )
 
 func init() {
@@ -91,7 +93,7 @@ func main() {
 		*publicIPv4, *publicIPv4Port,
 		*publicIPv6, *publicIPv6Port,
 		*statsIP, *statsPort,
-		*secret,
+		*secret, *adtag,
 	)
 	if err != nil {
 		usage(err.Error())
@@ -110,16 +112,23 @@ func main() {
 		zapcore.NewJSONEncoder(encoderCfg),
 		zapcore.Lock(os.Stderr),
 		atom,
-	)).Sugar()
+	))
+	zap.ReplaceGlobals(logger)
+	defer logger.Sync() // nolint: errcheck
 
-	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())
+	if conf.UseMiddleProxy() {
+		zap.S().Infow("Use middle proxy connection to Telegram")
+	} else {
+		zap.S().Infow("Use direct connection to Telegram")
+	}
+
+	go stats.Start(conf)
+
+	server := proxy.NewProxy(conf)
+	if err := server.Serve(); err != nil {
+		zap.S().Fatalw("Server stopped", "error", err)
 	}
 }