Direct proxy works

_antireplay/cache.go

@@ -0,0 +1,37 @@
+package antireplay
+
+import (
+	"github.com/allegro/bigcache"
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/config"
+)
+
+// Cache defines storage for obfuscated2 handshake frames.
+type Cache struct {
+	cache *bigcache.BigCache
+}
+
+func (a Cache) Add(frame []byte) {
+	a.cache.Set(string(frame), nil) // nolint: errcheck
+}
+
+func (a Cache) Has(frame []byte) bool {
+	_, err := a.cache.Get(string(frame))
+
+	return err == nil
+}
+
+func NewCache(config *config.Config) (Cache, error) {
+	cache, err := bigcache.NewBigCache(bigcache.Config{
+		Shards:           1024,
+		LifeWindow:       config.AntiReplayEvictionTime,
+		Hasher:           hasher{},
+		HardMaxCacheSize: config.AntiReplayMaxSize,
+	})
+	if err != nil {
+		return Cache{}, errors.Annotate(err, "Cannot make cache")
+	}
+
+	return Cache{cache}, nil
+}

newantireplay/hasher.go → _antireplay/hasher.go

@@ -1,4 +1,4 @@
-package newantireplay
+package antireplay
 
 import "github.com/cespare/xxhash"
 

_config/config.go

@@ -0,0 +1,224 @@
+package config
+
+import (
+	"bytes"
+	"encoding/hex"
+	"fmt"
+	"net"
+	"strconv"
+	"time"
+
+	"github.com/juju/errors"
+	statsd "gopkg.in/alexcesaro/statsd.v2"
+)
+
+// Config represents common configuration of mtg.
+type Config struct {
+	Debug      bool
+	Verbose    bool
+	SecureMode bool
+	SecureOnly bool
+
+	ReadBufferSize  int
+	WriteBufferSize int
+
+	BindPort       uint16
+	PublicIPv4Port uint16
+	PublicIPv6Port uint16
+	StatsPort      uint16
+
+	BindIP     net.IP
+	PublicIPv4 net.IP
+	PublicIPv6 net.IP
+	StatsIP    net.IP
+
+	AntiReplayMaxSize      int
+	AntiReplayEvictionTime time.Duration
+
+	StatsD struct {
+		Addr       net.Addr
+		Prefix     string
+		Tags       map[string]string
+		TagsFormat statsd.TagFormat
+		Enabled    bool
+	}
+	Prometheus struct {
+		Prefix string
+	}
+
+	Secret []byte
+	AdTag  []byte
+}
+
+// URLs contains links to the proxy (tg://, t.me) and their QR codes.
+type URLs struct {
+	TG        string `json:"tg_url"`
+	TMe       string `json:"tme_url"`
+	TGQRCode  string `json:"tg_qrcode"`
+	TMeQRCode string `json:"tme_qrcode"`
+}
+
+// IPURLs contains links to both ipv4 and ipv6 of the proxy.
+type IPURLs struct {
+	IPv4      URLs   `json:"ipv4"`
+	IPv6      URLs   `json:"ipv6"`
+	BotSecret string `json:"secret_for_mtproxybot"`
+}
+
+// BindAddr returns connection for this server to bind to.
+func (c *Config) BindAddr() string {
+	return getAddr(c.BindIP, c.BindPort)
+}
+
+// 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
+}
+
+// BotSecretString returns secret string which should work with MTProxybot.
+func (c *Config) BotSecretString() string {
+	return hex.EncodeToString(c.Secret)
+}
+
+// SecretString returns a secret in a form entered on the start of the
+// application.
+func (c *Config) SecretString() string {
+	secret := c.BotSecretString()
+	if c.SecureMode {
+		return "dd" + secret
+	}
+	return secret
+}
+
+// GetURLs returns configured IPURLs instance with links to this server.
+func (c *Config) GetURLs() IPURLs {
+	urls := IPURLs{}
+	secret := c.SecretString()
+	if c.PublicIPv4 != nil {
+		urls.IPv4 = getURLs(c.PublicIPv4, c.PublicIPv4Port, secret)
+	}
+	if c.PublicIPv6 != nil {
+		urls.IPv6 = getURLs(c.PublicIPv6, c.PublicIPv6Port, secret)
+	}
+	urls.BotSecret = c.BotSecretString()
+
+	return urls
+}
+
+func getAddr(host fmt.Stringer, port uint16) string {
+	return net.JoinHostPort(host.String(), strconv.Itoa(int(port)))
+}
+
+// NewConfig returns new configuration. If required, it manages and
+// fetches data from external sources. Parameters passed to this
+// function, should come from command line arguments.
+func NewConfig(debug, verbose bool, // nolint: gocyclo
+	writeBufferSize, readBufferSize uint32,
+	bindIP, publicIPv4, publicIPv6, statsIP net.IP,
+	bindPort, publicIPv4Port, publicIPv6Port, statsPort, statsdPort uint16,
+	statsdIP, statsdNetwork, statsdPrefix, statsdTagsFormat string,
+	statsdTags map[string]string, prometheusPrefix string,
+	secureOnly bool,
+	antiReplayMaxSize int, antiReplayEvictionTime time.Duration,
+	secret, adtag []byte) (*Config, error) {
+	secureMode := secureOnly
+	if bytes.HasPrefix(secret, []byte{0xdd}) && len(secret) == 17 {
+		secureMode = true
+		secret = bytes.TrimPrefix(secret, []byte{0xdd})
+	} else if len(secret) != 16 {
+		return nil, errors.New("Telegram demands secret of length 32")
+	}
+
+	var err error
+	if publicIPv4 == nil {
+		publicIPv4, err = getGlobalIPv4()
+		if err != nil {
+			publicIPv4 = nil
+		} else 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 = nil
+		} else if publicIPv6.To4() != nil {
+			return nil, errors.Errorf("IP %s is not IPv6", publicIPv6.String())
+		}
+	}
+	if publicIPv6Port == 0 {
+		publicIPv6Port = bindPort
+	}
+
+	if statsIP == nil {
+		statsIP = publicIPv4
+	}
+
+	conf := &Config{
+		Debug:                  debug,
+		Verbose:                verbose,
+		SecureOnly:             secureOnly,
+		BindIP:                 bindIP,
+		BindPort:               bindPort,
+		PublicIPv4:             publicIPv4,
+		PublicIPv4Port:         publicIPv4Port,
+		PublicIPv6:             publicIPv6,
+		PublicIPv6Port:         publicIPv6Port,
+		StatsIP:                statsIP,
+		StatsPort:              statsPort,
+		Secret:                 secret,
+		AdTag:                  adtag,
+		SecureMode:             secureMode,
+		ReadBufferSize:         int(readBufferSize),
+		WriteBufferSize:        int(writeBufferSize),
+		AntiReplayMaxSize:      antiReplayMaxSize,
+		AntiReplayEvictionTime: antiReplayEvictionTime,
+	}
+	conf.Prometheus.Prefix = prometheusPrefix
+
+	if statsdIP != "" {
+		conf.StatsD.Enabled = true
+		conf.StatsD.Prefix = statsdPrefix
+		conf.StatsD.Tags = statsdTags
+
+		var (
+			addr net.Addr
+			err  error
+		)
+		hostPort := net.JoinHostPort(statsdIP, strconv.Itoa(int(statsdPort)))
+		switch statsdNetwork {
+		case "tcp":
+			addr, err = net.ResolveTCPAddr("tcp", hostPort)
+		case "udp":
+			addr, err = net.ResolveUDPAddr("udp", hostPort)
+		default:
+			err = errors.Errorf("Unknown network %s", statsdNetwork)
+		}
+		if err != nil {
+			return nil, errors.Annotate(err, "Cannot resolve statsd address")
+		}
+		conf.StatsD.Addr = addr
+
+		switch statsdTagsFormat {
+		case "datadog":
+			conf.StatsD.TagsFormat = statsd.Datadog
+		case "influxdb":
+			conf.StatsD.TagsFormat = statsd.InfluxDB
+		case "":
+		default:
+			return nil, errors.Errorf("Unknown tags format %s", statsdTagsFormat)
+		}
+	}
+
+	return conf, nil
+}

newconfig/global_ips.go → _config/global_ips.go

@@ -1,43 +1,29 @@
-package newconfig
+package config
 
 import (
 	"context"
-	"io"
 	"io/ioutil"
 	"net"
 	"net/http"
 	"strings"
-	"time"
 
 	"github.com/juju/errors"
 )
 
-const (
-	ifconfigAddress = "https://ifconfig.co/ip"
-	ifconfigTimeout = 10 * time.Second
-)
+const ifconfigAddress = "https://ifconfig.co/ip"
 
 func getGlobalIPv4() (net.IP, error) {
-	ip, err := fetchIP("tcp4")
-	if err != nil || ip.To4() == nil {
-		return nil, errors.Annotate(err, "Cannot find public ipv4 address")
-	}
-	return ip, nil
+	return fetchIP("tcp4")
 }
 
 func getGlobalIPv6() (net.IP, error) {
-	ip, err := fetchIP("tcp6")
-	if err != nil || ip.To4() != nil {
-		return nil, errors.Annotate(err, "Cannot find public ipv6 address")
-	}
-	return ip, nil
+	return fetchIP("tcp6")
 }
 
 func fetchIP(network string) (net.IP, error) {
 	dialer := &net.Dialer{FallbackDelay: -1}
 	client := &http.Client{
-		Jar:     nil,
-		Timeout: ifconfigTimeout,
+		Jar: nil,
 		Transport: &http.Transport{
 			DialContext: func(ctx context.Context, _, addr string) (net.Conn, error) {
 				return dialer.DialContext(ctx, network, addr)
@@ -47,16 +33,13 @@ func fetchIP(network string) (net.IP, error) {
 
 	resp, err := client.Get(ifconfigAddress)
 	if err != nil {
-		if resp != nil {
-			io.Copy(ioutil.Discard, resp.Body) // nolint: errcheck
-		}
-		return nil, errors.Annotate(err, "Cannot perform a request")
+		return nil, err
 	}
 	defer resp.Body.Close() // nolint: errcheck
 
 	respDataBytes, err := ioutil.ReadAll(resp.Body)
 	if err != nil {
-		return nil, errors.Annotate(err, "Cannot read response body")
+		return nil, err
 	}
 	respData := strings.TrimSpace(string(respDataBytes))
 

newconfig/urls.go → _config/urls.go

@@ -1,42 +1,15 @@
-package newconfig
+package config
 
 import (
-	"encoding/hex"
+	"net"
 	"net/url"
+	"strconv"
 )
 
-type URLs struct {
-	TG        string `json:"tg_url"`
-	TMe       string `json:"tme_url"`
-	TGQRCode  string `json:"tg_qrcode"`
-	TMeQRCode string `json:"tme_qrcode"`
-}
-
-type IPURLs struct {
-	IPv4      URLs   `json:"ipv4"`
-	IPv6      URLs   `json:"ipv6"`
-	BotSecret string `json:"secret_for_mtproxybot"`
-}
-
-func GetURLs() (urls IPURLs) {
-	secret := ""
-	switch C.SecretMode {
-	case SecretModeSimple:
-		secret = hex.EncodeToString(C.Secret)
-	case SecretModeSecured:
-		secret = "dd" + hex.EncodeToString(C.Secret)
-	}
-
-	urls.IPv4 = makeURLs(&C.PublicIPv4Addr, secret)
-	urls.IPv6 = makeURLs(&C.PublicIPv6Addr, secret)
-	urls.BotSecret = secret
-
-	return urls
-}
-
-func makeURLs(addr *Addr, secret string) (urls URLs) {
+func getURLs(addr net.IP, port uint16, secret string) (urls URLs) {
 	values := url.Values{}
-	values.Set("address", addr.String())
+	values.Set("server", addr.String())
+	values.Set("port", strconv.Itoa(int(port)))
 	values.Set("secret", secret)
 
 	urls.TG = makeTGURL(values)

_obfuscated2/frame.go

@@ -0,0 +1,121 @@
+package obfuscated2
+
+import (
+	"bytes"
+	"crypto/rand"
+	"encoding/binary"
+	"io"
+
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/mtproto"
+)
+
+// [frameOffsetFirst:frameOffsetKey:frameOffsetIV:frameOffsetMagic:frameOffsetDC:frameOffsetEnd]
+const (
+	frameLenKey   = 32
+	frameLenIV    = 16
+	frameLenMagic = 4
+	frameLenDC    = 2
+
+	frameOffsetFirst = 8
+	frameOffsetKey   = frameOffsetFirst + frameLenKey
+	frameOffsetIV    = frameOffsetKey + frameLenIV
+	frameOffsetMagic = frameOffsetIV + frameLenMagic
+	frameOffsetDC    = frameOffsetMagic + frameLenDC
+
+	FrameLen = 64
+)
+
+// Frame represents handshake frame. Telegram sends 64 bytes of obfuscated2
+// initialization data first.
+// https://blog.susanka.eu/how-telegram-obfuscates-its-mtproto-traffic/
+type Frame []byte
+
+// Key returns AES encryption key.
+func (f Frame) Key() []byte {
+	return f[frameOffsetFirst:frameOffsetKey]
+}
+
+// IV returns AES encryption initialization vector
+func (f Frame) IV() []byte {
+	return f[frameOffsetKey:frameOffsetIV]
+}
+
+// Magic returns magic bytes from last 8 bytes of frame. Telegram checks
+// for values there. If after decryption magic is not as expected,
+// connection considered as failed.
+func (f Frame) Magic() []byte {
+	return f[frameOffsetIV:frameOffsetMagic]
+}
+
+// DC returns number of datacenter IP client wants to use.
+func (f Frame) DC() (n int16) {
+	buf := bytes.NewReader(f[frameOffsetMagic:frameOffsetDC])
+	if err := binary.Read(buf, binary.LittleEndian, &n); err != nil {
+		n = 1
+	}
+
+	return
+}
+
+// ConnectionType identifies connection type of the handshake frame.
+func (f Frame) ConnectionType() (mtproto.ConnectionType, error) {
+	return mtproto.ConnectionTagFromHandshake(f.Magic())
+}
+
+// Invert inverts frame for extracting encryption keys. Pkease check that link:
+// https://blog.susanka.eu/how-telegram-obfuscates-its-mtproto-traffic/
+func (f Frame) Invert() Frame {
+	reversed := make(Frame, FrameLen)
+	copy(reversed, f)
+
+	for i := 0; i < frameLenKey+frameLenIV; i++ {
+		reversed[frameOffsetFirst+i] = f[frameOffsetIV-1-i]
+	}
+
+	return reversed
+}
+
+// ExtractFrame extracts exact obfuscated2 handshake frame from given reader.
+func ExtractFrame(conn io.Reader) (Frame, error) {
+	frame := make(Frame, FrameLen)
+	buf := bytes.NewBuffer(frame)
+	buf.Reset()
+
+	if _, err := io.CopyN(buf, conn, FrameLen); err != nil {
+		return nil, errors.Annotate(err, "Cannot extract obfuscated header")
+	}
+	copy(frame, buf.Bytes())
+
+	return frame, nil
+}
+
+func generateFrame(connectionType mtproto.ConnectionType) Frame {
+	frame := make(Frame, FrameLen)
+
+	for {
+		if _, err := rand.Read(frame); err != nil {
+			continue
+		}
+		if frame[0] == 0xef {
+			continue
+		}
+
+		val := (uint32(frame[3]) << 24) | (uint32(frame[2]) << 16) | (uint32(frame[1]) << 8) | uint32(frame[0])
+		if val == 0x44414548 || val == 0x54534f50 || val == 0x20544547 || val == 0x4954504f || val == 0xeeeeeeee {
+			continue
+		}
+
+		val = (uint32(frame[7]) << 24) | (uint32(frame[6]) << 16) | (uint32(frame[5]) << 8) | uint32(frame[4])
+		if val == 0x00000000 {
+			continue
+		}
+
+		// error has to be checked before calling this function
+		tag, _ := connectionType.Tag() // nolint: errcheck, gosec
+		copy(frame.Magic(), tag)
+
+		return frame
+	}
+}

_proxy/proxy.go

@@ -0,0 +1,178 @@
+package proxy
+
+import (
+	"context"
+	"io"
+	"net"
+	"sync"
+
+	"github.com/gofrs/uuid"
+	"github.com/juju/errors"
+	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/antireplay"
+	"github.com/9seconds/mtg/client"
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/stats"
+	"github.com/9seconds/mtg/telegram"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+// Proxy is a core of this program.
+type Proxy struct {
+	antiReplayCache antireplay.Cache
+	clientInit      client.Init
+	tg              telegram.Telegram
+	conf            *config.Config
+}
+
+// Serve runs TCP proxy server.
+func (p *Proxy) Serve() error {
+	lsock, err := net.Listen("tcp", p.conf.BindAddr())
+	if err != nil {
+		return errors.Annotate(err, "Cannot create listen socket")
+	}
+
+	for {
+		if conn, err := lsock.Accept(); err != nil {
+			zap.S().Errorw("Cannot allocate incoming connection", "error", err)
+		} else {
+			go p.accept(conn)
+		}
+	}
+}
+
+func (p *Proxy) accept(conn net.Conn) {
+	connID := uuid.Must(uuid.NewV4()).String()
+	log := zap.S().With("connection_id", connID).Named("main")
+	ctx, cancel := context.WithCancel(context.Background())
+
+	defer func() {
+		cancel()
+		conn.Close() // nolint: errcheck, gosec
+
+		if err := recover(); err != nil {
+			stats.NewCrash()
+			log.Errorw("Crash of accept handler", "error", err)
+		}
+	}()
+
+	log.Infow("Client connected", "addr", conn.RemoteAddr())
+
+	clientConn, opts, err := p.clientInit(ctx, cancel, conn, connID, p.antiReplayCache, p.conf)
+	if err != nil {
+		log.Errorw("Cannot initialize client connection", "error", err)
+		return
+	}
+	defer clientConn.(io.Closer).Close() // nolint: errcheck
+
+	if p.conf.SecureOnly && opts.ConnectionType != mtproto.ConnectionTypeSecure {
+		log.Errorw("Proxy supports only secure connections", "connection_type", opts.ConnectionType)
+		return
+	}
+
+	stats.ClientConnected(opts.ConnectionType, clientConn.RemoteAddr())
+	defer stats.ClientDisconnected(opts.ConnectionType, clientConn.RemoteAddr())
+
+	serverConn, err := p.getTelegramConn(ctx, cancel, opts, connID)
+	if err != nil {
+		log.Errorw("Cannot initialize server connection", "error", err)
+		return
+	}
+	defer serverConn.(io.Closer).Close() // nolint: errcheck
+
+	go func() {
+		<-ctx.Done()
+		serverConn.(io.Closer).Close() // nolint: gosec
+		clientConn.(io.Closer).Close() // nolint: gosec
+	}()
+
+	wait := &sync.WaitGroup{}
+	wait.Add(2)
+
+	if p.conf.UseMiddleProxy() {
+		clientPacket := clientConn.(wrappers.PacketReadWriteCloser)
+		serverPacket := serverConn.(wrappers.PacketReadWriteCloser)
+		go p.middlePipe(clientPacket, serverPacket, wait, &opts.ReadHacks)
+		p.middlePipe(serverPacket, clientPacket, wait, &opts.WriteHacks)
+	} else {
+		clientStream := clientConn.(wrappers.StreamReadWriteCloser)
+		serverStream := serverConn.(wrappers.StreamReadWriteCloser)
+		go p.directPipe(clientStream, serverStream, wait, p.conf.ReadBufferSize)
+		p.directPipe(serverStream, clientStream, wait, p.conf.WriteBufferSize)
+	}
+
+	wait.Wait()
+
+	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(ctx, cancel, connID, opts)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot dial to Telegram")
+	}
+
+	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.PacketReadCloser, dst io.Writer, wait *sync.WaitGroup, hacks *mtproto.Hacks) {
+	defer wait.Done()
+
+	for {
+		hacks.SimpleAck = false
+		hacks.QuickAck = false
+
+		packet, err := src.Read()
+		if err != nil {
+			src.Logger().Warnw("Cannot read packet", "error", err)
+			return
+		}
+		if _, err = dst.Write(packet); err != nil {
+			src.Logger().Warnw("Cannot write packet", "error", err)
+			return
+		}
+	}
+}
+
+func (p *Proxy) directPipe(src wrappers.StreamReadCloser, dst io.Writer, wait *sync.WaitGroup, bufferSize int) {
+	defer wait.Done()
+
+	buffer := make([]byte, bufferSize)
+	if _, err := io.CopyBuffer(dst, src, buffer); err != nil {
+		src.Logger().Warnw("Cannot pump sockets", "error", err)
+	}
+}
+
+// NewProxy returns new proxy instance.
+func NewProxy(conf *config.Config) (*Proxy, error) {
+	var clientInit client.Init
+	var tg telegram.Telegram
+
+	cache, err := antireplay.NewCache(conf)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot make proxy")
+	}
+
+	if conf.UseMiddleProxy() {
+		clientInit = client.MiddleInit
+		tg = telegram.NewMiddleTelegram(conf)
+	} else {
+		clientInit = client.DirectInit
+		tg = telegram.NewDirectTelegram(conf)
+	}
+
+	return &Proxy{
+		antiReplayCache: cache,
+		conf:            conf,
+		clientInit:      clientInit,
+		tg:              tg,
+	}, nil
+}

_stats/stats.go

@@ -0,0 +1,175 @@
+package stats
+
+import (
+	"encoding/json"
+	"fmt"
+	"strconv"
+	"time"
+
+	humanize "github.com/dustin/go-humanize"
+
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
+)
+
+type uptime time.Time
+
+func (u uptime) MarshalJSON() ([]byte, error) {
+	duration := time.Since(time.Time(u))
+	value := map[string]string{
+		"seconds": strconv.Itoa(int(duration.Seconds())),
+		"human":   humanize.Time(time.Time(u)),
+	}
+
+	return json.Marshal(value)
+}
+
+type connectionType struct {
+	IPv6 uint32 `json:"ipv6"`
+	IPv4 uint32 `json:"ipv4"`
+}
+
+type baseConnections struct {
+	All          connectionType `json:"all"`
+	Abridged     connectionType `json:"abridged"`
+	Intermediate connectionType `json:"intermediate"`
+	Secure       connectionType `json:"secure"`
+}
+
+type connections struct {
+	baseConnections
+}
+
+func (c connections) MarshalJSON() ([]byte, error) {
+	c.All.IPv4 = c.Abridged.IPv4 + c.Intermediate.IPv4 + c.Secure.IPv4
+	c.All.IPv6 = c.Abridged.IPv6 + c.Intermediate.IPv6 + c.Secure.IPv6
+
+	return json.Marshal(c.baseConnections)
+}
+
+type traffic struct {
+	ingress uint64
+	egress  uint64
+}
+
+func (t *traffic) dumpValue(value uint64) map[string]interface{} {
+	return map[string]interface{}{
+		"bytes": value,
+		"human": humanize.Bytes(value),
+	}
+}
+
+func (t traffic) MarshalJSON() ([]byte, error) {
+	value := map[string]map[string]interface{}{
+		"ingress": t.dumpValue(t.ingress),
+		"egress":  t.dumpValue(t.egress),
+	}
+
+	return json.Marshal(value)
+}
+
+type speed struct {
+	ingress uint64
+	egress  uint64
+}
+
+func (s *speed) dumpValue(value uint64) map[string]interface{} {
+	return map[string]interface{}{
+		"bytes/s": value,
+		"human":   fmt.Sprintf("%s/s", humanize.Bytes(value)),
+	}
+}
+
+func (s speed) MarshalJSON() ([]byte, error) {
+	value := map[string]map[string]interface{}{
+		"ingress": s.dumpValue(s.ingress),
+		"egress":  s.dumpValue(s.egress),
+	}
+
+	return json.Marshal(value)
+}
+
+// Stats represents a statistics of the proxy.
+type Stats struct {
+	URLs        config.IPURLs `json:"urls"`
+	Connections connections   `json:"connections"`
+	Traffic     traffic       `json:"traffic"`
+	Speed       speed         `json:"speed"`
+	Uptime      uptime        `json:"uptime"`
+	Crashes     uint32        `json:"crashes"`
+
+	previousTraffic traffic
+}
+
+func (s *Stats) start() {
+	speedChan := time.Tick(time.Second)
+
+	for {
+		select {
+		case <-speedChan:
+			s.handleSpeed()
+		case event := <-trafficChan:
+			s.handleTraffic(event)
+		case event := <-connectionsChan:
+			s.handleConnection(event)
+		case getStatsChan := <-statsChan:
+			s.handleGetStats(getStatsChan)
+		case <-crashesChan:
+			s.handleCrash()
+		}
+	}
+}
+
+func (s *Stats) handleTraffic(evt trafficData) {
+	if evt.ingress {
+		s.Traffic.ingress += uint64(evt.traffic)
+	} else {
+		s.Traffic.egress += uint64(evt.traffic)
+	}
+}
+
+func (s *Stats) handleSpeed() {
+	s.Speed.ingress = s.Traffic.ingress - s.previousTraffic.ingress
+	s.Speed.egress = s.Traffic.egress - s.previousTraffic.egress
+	s.previousTraffic.ingress = s.Traffic.ingress
+	s.previousTraffic.egress = s.Traffic.egress
+}
+
+func (s *Stats) handleConnection(evt connectionData) {
+	var inc uint32 = 1
+	if !evt.connected {
+		inc = ^uint32(0)
+	}
+
+	var conn *connectionType
+	switch evt.connectionType {
+	case mtproto.ConnectionTypeAbridged:
+		conn = &s.Connections.Abridged
+	case mtproto.ConnectionTypeSecure:
+		conn = &s.Connections.Secure
+	default:
+		conn = &s.Connections.Intermediate
+	}
+
+	if evt.addr.IP.To4() != nil {
+		conn.IPv4 += inc
+	} else {
+		conn.IPv6 += inc
+	}
+}
+
+func (s *Stats) handleGetStats(getStatsChan chan<- Stats) {
+	getStatsChan <- *s
+}
+
+func (s *Stats) handleCrash() {
+	s.Crashes++
+}
+
+// NewStats creates a new instance of Stats structure.
+func NewStats(conf *config.Config) *Stats {
+	return &Stats{
+		URLs:   conf.GetURLs(),
+		Uptime: uptime(time.Now()),
+	}
+}

_telegram/direct.go

@@ -0,0 +1,79 @@
+package telegram
+
+import (
+	"context"
+	"net"
+
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/obfuscated2"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+const (
+	directV4DefaultIdx = 1
+	directV6DefaultIdx = 1
+)
+
+var (
+	directV4Addresses = map[int16][]string{
+		0: {"149.154.175.50:443"},
+		1: {"149.154.167.51:443"},
+		2: {"149.154.175.100:443"},
+		3: {"149.154.167.91:443"},
+		4: {"149.154.171.5:443"},
+	}
+	directV6Addresses = map[int16][]string{
+		0: {"[2001:b28:f23d:f001::a]:443"},
+		1: {"[2001:67c:04e8:f002::a]:443"},
+		2: {"[2001:b28:f23d:f003::a]:443"},
+		3: {"[2001:67c:04e8:f004::a]:443"},
+		4: {"[2001:b28:f23f:f005::a]:443"},
+	}
+)
+
+type directTelegram struct {
+	baseTelegram
+}
+
+func (t *directTelegram) Dial(ctx context.Context, cancel context.CancelFunc,
+	connID string, connOpts *mtproto.ConnectionOpts) (wrappers.StreamReadWriteCloser, error) {
+	dc := connOpts.DC
+	if dc < 0 {
+		dc = -dc
+	} else if dc == 0 {
+		dc = 1
+	}
+
+	return t.baseTelegram.dial(ctx, cancel, dc-1, connID, connOpts.ConnectionProto)
+}
+
+func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts,
+	conn wrappers.StreamReadWriteCloser) (wrappers.Wrap, error) {
+	obfs2, frame := obfuscated2.MakeTelegramObfuscated2Frame(connOpts)
+
+	if _, err := conn.Write(frame); err != nil {
+		return nil, errors.Annotate(err, "Cannot write hadnshake frame")
+	}
+
+	return wrappers.NewStreamCipher(conn, obfs2.Encryptor, obfs2.Decryptor), nil
+}
+
+// NewDirectTelegram returns Telegram instance which connects directly
+// to Telegram bypassing middleproxies.
+func NewDirectTelegram(conf *config.Config) Telegram {
+	return &directTelegram{
+		baseTelegram: baseTelegram{
+			dialer: tgDialer{
+				Dialer: net.Dialer{Timeout: telegramDialTimeout},
+				conf:   conf,
+			},
+			v4DefaultIdx: directV4DefaultIdx,
+			v6DefaultIdx: directV6DefaultIdx,
+			v4Addresses:  directV4Addresses,
+			v6Addresses:  directV6Addresses,
+		},
+	}
+}

_utils/reverse_bytes.go

@@ -0,0 +1,15 @@
+package utils
+
+// ReverseBytes is a common slice reverser.
+func ReverseBytes(data []byte) []byte {
+	dataLen := len(data)
+	rv := make([]byte, dataLen)
+
+	rv[dataLen/2] = data[dataLen/2]
+	for i := dataLen/2 - 1; i >= 0; i-- {
+		opp := dataLen - i - 1
+		rv[i], rv[opp] = data[opp], data[i]
+	}
+
+	return rv
+}

_utils/uint24.go

@@ -0,0 +1,15 @@
+package utils
+
+// Uint24 is a replacement for the absent Go uint24 data type.
+// This data type is little endian.
+type Uint24 [3]byte
+
+// ToUint24 converts number to Uint24.
+func ToUint24(number uint32) Uint24 {
+	return Uint24{byte(number), byte(number >> 8), byte(number >> 16)}
+}
+
+// FromUint24 converts Uint24 to number.
+func FromUint24(number Uint24) uint32 {
+	return uint32(number[0]) + (uint32(number[1]) << 8) + (uint32(number[2]) << 16)
+}

antireplay/cache.go

@@ -2,36 +2,29 @@ package antireplay
 
 import (
 	"github.com/allegro/bigcache"
-	"github.com/juju/errors"
 
 	"github.com/9seconds/mtg/config"
 )
 
-// Cache defines storage for obfuscated2 handshake frames.
-type Cache struct {
-	cache *bigcache.BigCache
-}
+var cache *bigcache.BigCache
 
-func (a Cache) Add(frame []byte) {
-	a.cache.Set(string(frame), nil) // nolint: errcheck
+func Add(data []byte) {
+	cache.Set(string(data), nil) // nolint: errcheck
 }
 
-func (a Cache) Has(frame []byte) bool {
-	_, err := a.cache.Get(string(frame))
-
+func Has(data []byte) bool {
+	_, err := cache.Get(string(data))
 	return err == nil
 }
 
-func NewCache(config *config.Config) (Cache, error) {
-	cache, err := bigcache.NewBigCache(bigcache.Config{
+func Init() error {
+	c, err := bigcache.NewBigCache(bigcache.Config{
 		Shards:           1024,
-		LifeWindow:       config.AntiReplayEvictionTime,
+		LifeWindow:       config.C.AntiReplay.EvictionTime,
 		Hasher:           hasher{},
-		HardMaxCacheSize: config.AntiReplayMaxSize,
+		HardMaxCacheSize: config.C.AntiReplay.MaxSize,
 	})
-	if err != nil {
-		return Cache{}, errors.Annotate(err, "Cannot make cache")
-	}
+	cache = c
 
-	return Cache{cache}, nil
+	return err
 }

newcli/generate.go → cli/generate.go

@@ -1,14 +1,14 @@
-package newcli
+package cli
 
 import (
 	"crypto/rand"
 	"encoding/hex"
 
-	"github.com/9seconds/mtg/newconfig"
+	"github.com/9seconds/mtg/config"
 )
 
 func Generate(secretType string) {
-	data := make([]byte, newconfig.SimpleSecretLength)
+	data := make([]byte, config.SimpleSecretLength)
 	if _, err := rand.Read(data); err != nil {
 		panic(err)
 	}

cli/proxy.go

@@ -0,0 +1,86 @@
+package cli
+
+import (
+	"net"
+	"os"
+	"time"
+
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+
+	"github.com/9seconds/mtg/antireplay"
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/ntp"
+	"github.com/9seconds/mtg/obfuscated2"
+	"github.com/9seconds/mtg/proxy"
+	"github.com/9seconds/mtg/stats"
+	"github.com/9seconds/mtg/telegram"
+)
+
+func Proxy() error {
+	atom := zap.NewAtomicLevel()
+	switch {
+	case config.C.Debug:
+		atom.SetLevel(zapcore.DebugLevel)
+	case config.C.Verbose:
+		atom.SetLevel(zapcore.InfoLevel)
+	default:
+		atom.SetLevel(zapcore.ErrorLevel)
+	}
+
+	encoderCfg := zap.NewProductionEncoderConfig()
+	logger := zap.New(zapcore.NewCore(
+		zapcore.NewJSONEncoder(encoderCfg),
+		zapcore.Lock(os.Stderr),
+		atom,
+	))
+	zap.ReplaceGlobals(logger)
+	defer logger.Sync() // nolint: errcheck
+
+	if err := config.InitPublicAddress(); err != nil {
+		Fatal(err.Error())
+	}
+	zap.S().Debugw("Configuration", "config", config.C)
+
+	if len(config.C.AdTag) > 0 {
+		zap.S().Infow("Use middle proxy connection to Telegram")
+		diff, err := ntp.Fetch()
+		if err != nil {
+			Fatal("Cannot fetch time data from NTP")
+		}
+		if diff > time.Second {
+			Fatal("Your local time is skewed and drift is bigger than a second. Please sync your time.")
+		}
+		go ntp.AutoUpdate()
+	} else {
+		zap.S().Infow("Use direct connection to Telegram")
+	}
+
+	PrintJSONStdout(config.GetURLs())
+
+	if err := antireplay.Init(); err != nil {
+		Fatal(err.Error())
+	}
+	if err := stats.Init(); err != nil {
+		Fatal(err.Error())
+	}
+	proxyListener, err := net.Listen("tcp", config.C.ListenAddr.String())
+	if err != nil {
+		Fatal(err.Error())
+	}
+
+	app := &proxy.Proxy{
+		Logger: zap.S().Named("proxy"),
+	}
+	if len(config.C.AdTag) == 0 {
+		app.TelegramProtocolMaker = obfuscated2.MakeTelegramProtocol
+		app.TelegramDialer = telegram.NewDirectTelegram()
+	}
+	if config.C.SecretMode != config.SecretModeTLS {
+		app.ClientProtocolMaker = obfuscated2.MakeClientProtocol
+	}
+
+	app.Serve(proxyListener)
+
+	return nil
+}

newcli/utils.go → cli/utils.go

@@ -1,4 +1,4 @@
-package newcli
+package cli
 
 import (
 	"encoding/json"

config/config.go

@@ -2,223 +2,257 @@ package config
 
 import (
 	"bytes"
-	"encoding/hex"
-	"fmt"
+	"encoding/json"
 	"net"
 	"strconv"
 	"time"
 
 	"github.com/juju/errors"
+	"go.uber.org/zap"
 	statsd "gopkg.in/alexcesaro/statsd.v2"
 )
 
-// Config represents common configuration of mtg.
-type Config struct {
-	Debug      bool
-	Verbose    bool
-	SecureMode bool
-	SecureOnly bool
-
-	ReadBufferSize  int
-	WriteBufferSize int
-
-	BindPort       uint16
-	PublicIPv4Port uint16
-	PublicIPv6Port uint16
-	StatsPort      uint16
-
-	BindIP     net.IP
-	PublicIPv4 net.IP
-	PublicIPv6 net.IP
-	StatsIP    net.IP
-
-	AntiReplayMaxSize      int
-	AntiReplayEvictionTime time.Duration
-
-	StatsD struct {
-		Addr       net.Addr
-		Prefix     string
-		Tags       map[string]string
-		TagsFormat statsd.TagFormat
-		Enabled    bool
-	}
-	Prometheus struct {
-		Prefix string
-	}
+type SecretMode uint8
 
-	Secret []byte
-	AdTag  []byte
+func (s SecretMode) String() string {
+	switch s {
+	case SecretModeSimple:
+		return "simple"
+	case SecretModeSecured:
+		return "secured"
+	}
+	return "tls"
 }
 
-// URLs contains links to the proxy (tg://, t.me) and their QR codes.
-type URLs struct {
-	TG        string `json:"tg_url"`
-	TMe       string `json:"tme_url"`
-	TGQRCode  string `json:"tg_qrcode"`
-	TMeQRCode string `json:"tme_qrcode"`
+const (
+	SecretModeSimple SecretMode = iota
+	SecretModeSecured
+	SecretModeTLS
+)
+
+const SimpleSecretLength = 16
+
+type OptionType uint8
+
+const (
+	OptionTypeDebug OptionType = iota
+	OptionTypeVerbose
+
+	OptionTypeBindIP
+	OptionTypeBindPort
+	OptionTypePublicIPv4
+	OptionTypePublicIPv4Port
+	OptionTypePublicIPv6
+	OptionTypePublicIPv6Port
+	OptionTypeStatsIP
+	OptionTypeStatsPort
+
+	OptionTypeStatsdIP
+	OptionTypeStatsdPort
+	OptionTypeStatsdNetwork
+	OptionTypeStatsdPrefix
+	OptionTypeStatsdTagsFormat
+	OptionTypeStatsdTags
+	OptionTypePrometheusPrefix
+
+	OptionTypeWriteBufferSize
+	OptionTypeReadBufferSize
+
+	OptionTypeAntiReplayMaxSize
+	OptionTypeAntiReplayEvictionTime
+
+	OptionTypeSecret
+	OptionTypeAdtag
+)
+
+type BufferSize struct {
+	Read  int `json:"read"`
+	Write int `json:"write"`
 }
 
-// IPURLs contains links to both ipv4 and ipv6 of the proxy.
-type IPURLs struct {
-	IPv4      URLs   `json:"ipv4"`
-	IPv6      URLs   `json:"ipv6"`
-	BotSecret string `json:"secret_for_mtproxybot"`
+type AntiReplay struct {
+	MaxSize      int           `json:"max_size"`
+	EvictionTime time.Duration `json:"duration"`
 }
 
-// BindAddr returns connection for this server to bind to.
-func (c *Config) BindAddr() string {
-	return getAddr(c.BindIP, c.BindPort)
+type Stats struct {
+	Prefix string `json:"prefix"`
 }
 
-// StatAddr returns connection string to the stats API.
-func (c *Config) StatAddr() string {
-	return getAddr(c.StatsIP, c.StatsPort)
+type StatsdStats struct {
+	Stats
+
+	Addr       Addr              `json:"addr"`
+	Tags       map[string]string `json:"tags"`
+	TagsFormat statsd.TagFormat  `json:"format"`
 }
 
-// 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
+type PrometheusStats struct {
+	Stats
 }
 
-// BotSecretString returns secret string which should work with MTProxybot.
-func (c *Config) BotSecretString() string {
-	return hex.EncodeToString(c.Secret)
+type Addr struct {
+	IP   net.IP `json:"ip"`
+	Port int    `json:"port"`
+	net  string
 }
 
-// SecretString returns a secret in a form entered on the start of the
-// application.
-func (c *Config) SecretString() string {
-	secret := c.BotSecretString()
-	if c.SecureMode {
-		return "dd" + secret
+func (a Addr) Network() string {
+	if a.net == "" {
+		return "tcp"
 	}
-	return secret
+	return a.net
 }
 
-// GetURLs returns configured IPURLs instance with links to this server.
-func (c *Config) GetURLs() IPURLs {
-	urls := IPURLs{}
-	secret := c.SecretString()
-	if c.PublicIPv4 != nil {
-		urls.IPv4 = getURLs(c.PublicIPv4, c.PublicIPv4Port, secret)
-	}
-	if c.PublicIPv6 != nil {
-		urls.IPv6 = getURLs(c.PublicIPv6, c.PublicIPv6Port, secret)
+func (a Addr) String() string {
+	return net.JoinHostPort(a.IP.String(), strconv.Itoa(a.Port))
+}
+
+func (a Addr) MarshalJSON() ([]byte, error) {
+	data := map[string]string{
+		"network": a.Network(),
+		"addr":    a.String(),
 	}
-	urls.BotSecret = c.BotSecretString()
+	return json.Marshal(data)
+}
+
+type Config struct {
+	BufferSize BufferSize `json:"buffer_size"`
+	AntiReplay AntiReplay `json:"anti_replay"`
 
-	return urls
+	ListenAddr     Addr `json:"listen_addr"`
+	PublicIPv4Addr Addr `json:"public_ipv4_addr"`
+	PublicIPv6Addr Addr `json:"public_ipv6_addr"`
+	StatsAddr      Addr `json:"stats_addr"`
+
+	StatsdStats     StatsdStats     `json:"stats_statsd"`
+	PrometheusStats PrometheusStats `json:"stats_prometheus"`
+
+	Debug      bool       `json:"debug"`
+	Verbose    bool       `json:"verbose"`
+	SecretMode SecretMode `json:"secret_mode"`
+	Secret     []byte     `json:"secret"`
+	AdTag      []byte     `json:"adtag"`
 }
 
-func getAddr(host fmt.Stringer, port uint16) string {
-	return net.JoinHostPort(host.String(), strconv.Itoa(int(port)))
+func (c Config) String() string {
+	data, _ := json.Marshal(c)
+	return string(data)
 }
 
-// NewConfig returns new configuration. If required, it manages and
-// fetches data from external sources. Parameters passed to this
-// function, should come from command line arguments.
-func NewConfig(debug, verbose bool, // nolint: gocyclo
-	writeBufferSize, readBufferSize uint32,
-	bindIP, publicIPv4, publicIPv6, statsIP net.IP,
-	bindPort, publicIPv4Port, publicIPv6Port, statsPort, statsdPort uint16,
-	statsdIP, statsdNetwork, statsdPrefix, statsdTagsFormat string,
-	statsdTags map[string]string, prometheusPrefix string,
-	secureOnly bool,
-	antiReplayMaxSize int, antiReplayEvictionTime time.Duration,
-	secret, adtag []byte) (*Config, error) {
-	secureMode := secureOnly
-	if bytes.HasPrefix(secret, []byte{0xdd}) && len(secret) == 17 {
-		secureMode = true
-		secret = bytes.TrimPrefix(secret, []byte{0xdd})
-	} else if len(secret) != 16 {
-		return nil, errors.New("Telegram demands secret of length 32")
-	}
+type ConfigOpt struct {
+	Option OptionType
+	Value  interface{}
+}
 
-	var err error
-	if publicIPv4 == nil {
-		publicIPv4, err = getGlobalIPv4()
-		if err != nil {
-			publicIPv4 = nil
-		} else if publicIPv4.To4() == nil {
-			return nil, errors.Errorf("IP %s is not IPv4", publicIPv4.String())
-		}
-	}
-	if publicIPv4Port == 0 {
-		publicIPv4Port = bindPort
-	}
+var C = Config{}
 
-	if publicIPv6 == nil {
-		publicIPv6, err = getGlobalIPv6()
-		if err != nil {
-			publicIPv6 = nil
-		} else if publicIPv6.To4() != nil {
-			return nil, errors.Errorf("IP %s is not IPv6", publicIPv6.String())
+func Init(options ...ConfigOpt) error { // nolint: gocyclo
+	for _, opt := range options {
+		switch opt.Option {
+		case OptionTypeDebug:
+			C.Debug = opt.Value.(bool)
+		case OptionTypeVerbose:
+			C.Verbose = opt.Value.(bool)
+		case OptionTypeBindIP:
+			C.ListenAddr.IP = opt.Value.(net.IP)
+		case OptionTypeBindPort:
+			C.ListenAddr.Port = int(opt.Value.(uint16))
+		case OptionTypePublicIPv4:
+			C.PublicIPv4Addr.IP = opt.Value.(net.IP)
+		case OptionTypePublicIPv4Port:
+			C.PublicIPv4Addr.Port = int(opt.Value.(uint16))
+		case OptionTypePublicIPv6:
+			C.PublicIPv6Addr.IP = opt.Value.(net.IP)
+		case OptionTypePublicIPv6Port:
+			C.PublicIPv6Addr.Port = int(opt.Value.(uint16))
+		case OptionTypeStatsIP:
+			C.StatsAddr.IP = opt.Value.(net.IP)
+		case OptionTypeStatsPort:
+			C.StatsAddr.Port = int(opt.Value.(uint16))
+		case OptionTypeStatsdIP:
+			C.StatsdStats.Addr.IP = opt.Value.(net.IP)
+		case OptionTypeStatsdPort:
+			C.StatsdStats.Addr.Port = int(opt.Value.(uint16))
+		case OptionTypeStatsdNetwork:
+			C.StatsdStats.Addr.net = opt.Value.(string)
+		case OptionTypeStatsdPrefix:
+			C.StatsdStats.Prefix = opt.Value.(string)
+		case OptionTypeStatsdTagsFormat:
+			value := opt.Value.(string)
+			switch value {
+			case "datadog":
+				C.StatsdStats.TagsFormat = statsd.Datadog
+			case "influxdb":
+				C.StatsdStats.TagsFormat = statsd.InfluxDB
+			default:
+				return errors.Errorf("Incorrect statsd tag %s", value)
+			}
+		case OptionTypeStatsdTags:
+			C.StatsdStats.Tags = opt.Value.(map[string]string)
+		case OptionTypePrometheusPrefix:
+			C.PrometheusStats.Prefix = opt.Value.(string)
+		case OptionTypeWriteBufferSize:
+			C.BufferSize.Write = int(opt.Value.(uint32))
+		case OptionTypeReadBufferSize:
+			C.BufferSize.Read = int(opt.Value.(uint32))
+		case OptionTypeAntiReplayMaxSize:
+			C.AntiReplay.MaxSize = opt.Value.(int)
+		case OptionTypeAntiReplayEvictionTime:
+			C.AntiReplay.EvictionTime = opt.Value.(time.Duration)
+		case OptionTypeSecret:
+			C.Secret = opt.Value.([]byte)
+		case OptionTypeAdtag:
+			C.AdTag = opt.Value.([]byte)
+		default:
+			return errors.Errorf("Unknown tag %v", opt.Option)
 		}
 	}
-	if publicIPv6Port == 0 {
-		publicIPv6Port = bindPort
-	}
 
-	if statsIP == nil {
-		statsIP = publicIPv4
+	switch {
+	case len(C.Secret) == 1+SimpleSecretLength && bytes.HasPrefix(C.Secret, []byte{0xdd}):
+		C.SecretMode = SecretModeSecured
+		C.Secret = bytes.TrimPrefix(C.Secret, []byte{0xdd})
+	case len(C.Secret) == SimpleSecretLength:
+		C.SecretMode = SecretModeSimple
+	default:
+		return errors.New("Incorrect secret")
 	}
 
-	conf := &Config{
-		Debug:                  debug,
-		Verbose:                verbose,
-		SecureOnly:             secureOnly,
-		BindIP:                 bindIP,
-		BindPort:               bindPort,
-		PublicIPv4:             publicIPv4,
-		PublicIPv4Port:         publicIPv4Port,
-		PublicIPv6:             publicIPv6,
-		PublicIPv6Port:         publicIPv6Port,
-		StatsIP:                statsIP,
-		StatsPort:              statsPort,
-		Secret:                 secret,
-		AdTag:                  adtag,
-		SecureMode:             secureMode,
-		ReadBufferSize:         int(readBufferSize),
-		WriteBufferSize:        int(writeBufferSize),
-		AntiReplayMaxSize:      antiReplayMaxSize,
-		AntiReplayEvictionTime: antiReplayEvictionTime,
+	return nil
+}
+
+func InitPublicAddress() error {
+	if C.PublicIPv4Addr.Port == 0 {
+		C.PublicIPv4Addr.Port = C.ListenAddr.Port
 	}
-	conf.Prometheus.Prefix = prometheusPrefix
-
-	if statsdIP != "" {
-		conf.StatsD.Enabled = true
-		conf.StatsD.Prefix = statsdPrefix
-		conf.StatsD.Tags = statsdTags
-
-		var (
-			addr net.Addr
-			err  error
-		)
-		hostPort := net.JoinHostPort(statsdIP, strconv.Itoa(int(statsdPort)))
-		switch statsdNetwork {
-		case "tcp":
-			addr, err = net.ResolveTCPAddr("tcp", hostPort)
-		case "udp":
-			addr, err = net.ResolveUDPAddr("udp", hostPort)
-		default:
-			err = errors.Errorf("Unknown network %s", statsdNetwork)
-		}
+	if C.PublicIPv6Addr.Port == 0 {
+		C.PublicIPv6Addr.Port = C.ListenAddr.Port
+	}
+
+	foundAddress := C.PublicIPv4Addr.IP != nil || C.PublicIPv6Addr.IP != nil
+	if C.PublicIPv4Addr.IP == nil {
+		ip, err := getGlobalIPv4()
 		if err != nil {
-			return nil, errors.Annotate(err, "Cannot resolve statsd address")
+			zap.S().Warnw("Cannot resolve public address", "error", err)
+		} else {
+			C.PublicIPv4Addr.IP = ip
+			foundAddress = true
 		}
-		conf.StatsD.Addr = addr
-
-		switch statsdTagsFormat {
-		case "datadog":
-			conf.StatsD.TagsFormat = statsd.Datadog
-		case "influxdb":
-			conf.StatsD.TagsFormat = statsd.InfluxDB
-		case "":
-		default:
-			return nil, errors.Errorf("Unknown tags format %s", statsdTagsFormat)
+	}
+	if C.PublicIPv6Addr.IP == nil {
+		ip, err := getGlobalIPv6()
+		if err != nil {
+			zap.S().Warnw("Cannot resolve public address", "error", err)
+		} else {
+			C.PublicIPv6Addr.IP = ip
+			foundAddress = true
 		}
 	}
 
-	return conf, nil
+	if !foundAddress {
+		return errors.New("Cannot resolve any public address")
+	}
+
+	return nil
 }

config/global_ips.go

@@ -2,28 +2,42 @@ package config
 
 import (
 	"context"
+	"io"
 	"io/ioutil"
 	"net"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/juju/errors"
 )
 
-const ifconfigAddress = "https://ifconfig.co/ip"
+const (
+	ifconfigAddress = "https://ifconfig.co/ip"
+	ifconfigTimeout = 10 * time.Second
+)
 
 func getGlobalIPv4() (net.IP, error) {
-	return fetchIP("tcp4")
+	ip, err := fetchIP("tcp4")
+	if err != nil || ip.To4() == nil {
+		return nil, errors.Annotate(err, "Cannot find public ipv4 address")
+	}
+	return ip, nil
 }
 
 func getGlobalIPv6() (net.IP, error) {
-	return fetchIP("tcp6")
+	ip, err := fetchIP("tcp6")
+	if err != nil || ip.To4() != nil {
+		return nil, errors.Annotate(err, "Cannot find public ipv6 address")
+	}
+	return ip, nil
 }
 
 func fetchIP(network string) (net.IP, error) {
 	dialer := &net.Dialer{FallbackDelay: -1}
 	client := &http.Client{
-		Jar: nil,
+		Jar:     nil,
+		Timeout: ifconfigTimeout,
 		Transport: &http.Transport{
 			DialContext: func(ctx context.Context, _, addr string) (net.Conn, error) {
 				return dialer.DialContext(ctx, network, addr)
@@ -33,13 +47,16 @@ func fetchIP(network string) (net.IP, error) {
 
 	resp, err := client.Get(ifconfigAddress)
 	if err != nil {
-		return nil, err
+		if resp != nil {
+			io.Copy(ioutil.Discard, resp.Body) // nolint: errcheck
+		}
+		return nil, errors.Annotate(err, "Cannot perform a request")
 	}
 	defer resp.Body.Close() // nolint: errcheck
 
 	respDataBytes, err := ioutil.ReadAll(resp.Body)
 	if err != nil {
-		return nil, err
+		return nil, errors.Annotate(err, "Cannot read response body")
 	}
 	respData := strings.TrimSpace(string(respDataBytes))
 

config/urls.go

@@ -1,15 +1,42 @@
 package config
 
 import (
-	"net"
+	"encoding/hex"
 	"net/url"
-	"strconv"
 )
 
-func getURLs(addr net.IP, port uint16, secret string) (urls URLs) {
+type URLs struct {
+	TG        string `json:"tg_url"`
+	TMe       string `json:"tme_url"`
+	TGQRCode  string `json:"tg_qrcode"`
+	TMeQRCode string `json:"tme_qrcode"`
+}
+
+type IPURLs struct {
+	IPv4      URLs   `json:"ipv4"`
+	IPv6      URLs   `json:"ipv6"`
+	BotSecret string `json:"secret_for_mtproxybot"`
+}
+
+func GetURLs() (urls IPURLs) {
+	secret := ""
+	switch C.SecretMode {
+	case SecretModeSimple:
+		secret = hex.EncodeToString(C.Secret)
+	case SecretModeSecured:
+		secret = "dd" + hex.EncodeToString(C.Secret)
+	}
+
+	urls.IPv4 = makeURLs(&C.PublicIPv4Addr, secret)
+	urls.IPv6 = makeURLs(&C.PublicIPv6Addr, secret)
+	urls.BotSecret = secret
+
+	return urls
+}
+
+func makeURLs(addr *Addr, secret string) (urls URLs) {
 	values := url.Values{}
-	values.Set("server", addr.String())
-	values.Set("port", strconv.Itoa(int(port)))
+	values.Set("address", addr.String())
 	values.Set("secret", secret)
 
 	urls.TG = makeTGURL(values)

conntypes/dc.go

@@ -0,0 +1,5 @@
+package conntypes
+
+type DC int16
+
+const DCDefaultIdx DC = 1

newprotocol/connection_protocol.go → conntypes/protocol.go

@@ -1,4 +1,4 @@
-package newprotocol
+package conntypes
 
 type ConnectionProtocol uint8
 

newprotocol/connection_type.go → conntypes/type.go

@@ -1,4 +1,4 @@
-package newprotocol
+package conntypes
 
 type ConnectionType uint8
 

main.go

@@ -9,8 +9,8 @@ import (
 	"github.com/juju/errors"
 	kingpin "gopkg.in/alecthomas/kingpin.v2"
 
-	"github.com/9seconds/mtg/newcli"
-	"github.com/9seconds/mtg/newconfig"
+	"github.com/9seconds/mtg/cli"
+	"github.com/9seconds/mtg/config"
 )
 
 var version = "dev" // this has to be set by build ld flags
@@ -144,45 +144,45 @@ func main() {
 	app.HelpFlag.Short('h')
 
 	if err := setRLimit(); err != nil {
-		newcli.Fatal(err.Error())
+		cli.Fatal(err.Error())
 	}
 
 	switch kingpin.MustParse(app.Parse(os.Args[1:])) {
 	case generateSecretCommand.FullCommand():
-		newcli.Generate(*generateSecretType)
+		cli.Generate(*generateSecretType)
 
 	case proxyCommand.FullCommand():
-		err := newconfig.Init(
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeDebug, Value: *proxyDebug},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeVerbose, Value: *proxyVerbose},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeBindIP, Value: *proxyBindIP},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeBindPort, Value: *proxyBindPort},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypePublicIPv4, Value: *proxyPublicIPv4},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypePublicIPv4Port, Value: *proxyPublicIPv4Port},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypePublicIPv6, Value: *proxyPublicIPv6},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypePublicIPv6Port, Value: *proxyPublicIPv6Port},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeStatsIP, Value: *proxyStatsIP},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeStatsPort, Value: *proxyStatsPort},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeStatsdIP, Value: *proxyStatsdIP},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeStatsdPort, Value: *proxyStatsdPort},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeStatsdNetwork, Value: *proxyStatsdNetwork},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeStatsdPrefix, Value: *proxyStatsdPrefix},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeStatsdTagsFormat, Value: *proxyStatsdTagsFormat},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeStatsdTags, Value: *proxyStatsdTags},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypePrometheusPrefix, Value: *proxyPrometheusPrefix},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeWriteBufferSize, Value: *proxyWriteBufferSize},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeReadBufferSize, Value: *proxyReadBufferSize},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeAntiReplayMaxSize, Value: *proxyAntiReplayMaxSize},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeAntiReplayEvictionTime, Value: *proxyAntiReplayEvictionTime},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeSecret, Value: *proxySecret},
-			newconfig.ConfigOpt{Option: newconfig.OptionTypeAdtag, Value: *proxyAdtag},
+		err := config.Init(
+			config.ConfigOpt{Option: config.OptionTypeDebug, Value: *proxyDebug},
+			config.ConfigOpt{Option: config.OptionTypeVerbose, Value: *proxyVerbose},
+			config.ConfigOpt{Option: config.OptionTypeBindIP, Value: *proxyBindIP},
+			config.ConfigOpt{Option: config.OptionTypeBindPort, Value: *proxyBindPort},
+			config.ConfigOpt{Option: config.OptionTypePublicIPv4, Value: *proxyPublicIPv4},
+			config.ConfigOpt{Option: config.OptionTypePublicIPv4Port, Value: *proxyPublicIPv4Port},
+			config.ConfigOpt{Option: config.OptionTypePublicIPv6, Value: *proxyPublicIPv6},
+			config.ConfigOpt{Option: config.OptionTypePublicIPv6Port, Value: *proxyPublicIPv6Port},
+			config.ConfigOpt{Option: config.OptionTypeStatsIP, Value: *proxyStatsIP},
+			config.ConfigOpt{Option: config.OptionTypeStatsPort, Value: *proxyStatsPort},
+			config.ConfigOpt{Option: config.OptionTypeStatsdIP, Value: *proxyStatsdIP},
+			config.ConfigOpt{Option: config.OptionTypeStatsdPort, Value: *proxyStatsdPort},
+			config.ConfigOpt{Option: config.OptionTypeStatsdNetwork, Value: *proxyStatsdNetwork},
+			config.ConfigOpt{Option: config.OptionTypeStatsdPrefix, Value: *proxyStatsdPrefix},
+			config.ConfigOpt{Option: config.OptionTypeStatsdTagsFormat, Value: *proxyStatsdTagsFormat},
+			config.ConfigOpt{Option: config.OptionTypeStatsdTags, Value: *proxyStatsdTags},
+			config.ConfigOpt{Option: config.OptionTypePrometheusPrefix, Value: *proxyPrometheusPrefix},
+			config.ConfigOpt{Option: config.OptionTypeWriteBufferSize, Value: *proxyWriteBufferSize},
+			config.ConfigOpt{Option: config.OptionTypeReadBufferSize, Value: *proxyReadBufferSize},
+			config.ConfigOpt{Option: config.OptionTypeAntiReplayMaxSize, Value: *proxyAntiReplayMaxSize},
+			config.ConfigOpt{Option: config.OptionTypeAntiReplayEvictionTime, Value: *proxyAntiReplayEvictionTime},
+			config.ConfigOpt{Option: config.OptionTypeSecret, Value: *proxySecret},
+			config.ConfigOpt{Option: config.OptionTypeAdtag, Value: *proxyAdtag},
 		)
 		if err != nil {
-			newcli.Fatal(err.Error())
+			cli.Fatal(err.Error())
 		}
 
-		if err := newcli.Proxy(); err != nil {
-			newcli.Fatal(err.Error())
+		if err := cli.Proxy(); err != nil {
+			cli.Fatal(err.Error())
 		}
 	}
 }

newantireplay/cache.go

@@ -1,32 +0,0 @@
-package newantireplay
-
-import (
-	"github.com/allegro/bigcache"
-
-	"github.com/9seconds/mtg/newconfig"
-)
-
-var cache *bigcache.BigCache
-
-func Add(data []byte) {
-	cache.Set(string(data), nil)
-}
-
-func Has(data []byte) bool {
-	_, err := cache.Get(string(data))
-	return err == nil
-}
-
-func Init() {
-	c, err := bigcache.NewBigCache(bigcache.Config{
-		Shards:           1024,
-		LifeWindow:       newconfig.C.AntiReplay.EvictionTime,
-		Hasher:           hasher{},
-		HardMaxCacheSize: newconfig.C.AntiReplay.MaxSize,
-	})
-	if err != nil {
-		panic(err)
-	}
-
-	cache = c
-}

newcli/proxy.go

@@ -1,61 +0,0 @@
-package newcli
-
-import (
-	"os"
-	"time"
-
-	"go.uber.org/zap"
-	"go.uber.org/zap/zapcore"
-
-	"github.com/9seconds/mtg/newconfig"
-	"github.com/9seconds/mtg/newstats"
-	"github.com/9seconds/mtg/ntp"
-)
-
-func Proxy() error {
-	atom := zap.NewAtomicLevel()
-	switch {
-	case newconfig.C.Debug:
-		atom.SetLevel(zapcore.DebugLevel)
-	case newconfig.C.Verbose:
-		atom.SetLevel(zapcore.InfoLevel)
-	default:
-		atom.SetLevel(zapcore.ErrorLevel)
-	}
-
-	encoderCfg := zap.NewProductionEncoderConfig()
-	logger := zap.New(zapcore.NewCore(
-		zapcore.NewJSONEncoder(encoderCfg),
-		zapcore.Lock(os.Stderr),
-		atom,
-	))
-	zap.ReplaceGlobals(logger)
-	defer logger.Sync() // nolint: errcheck
-
-	if err := newconfig.InitPublicAddress(); err != nil {
-		Fatal(err.Error())
-	}
-	zap.S().Debugw("Configuration", "config", newconfig.C)
-
-	if len(newconfig.C.AdTag) > 0 {
-		zap.S().Infow("Use middle proxy connection to Telegram")
-		diff, err := ntp.Fetch()
-		if err != nil {
-			Fatal("Cannot fetch time data from NTP")
-		}
-		if diff > time.Second {
-			Fatal("Your local time is skewed and drift is bigger than a second. Please sync your time.")
-		}
-		go ntp.AutoUpdate()
-	} else {
-		zap.S().Infow("Use direct connection to Telegram")
-	}
-
-	PrintJSONStdout(newconfig.GetURLs())
-
-	if err := newstats.Init(); err != nil {
-		Fatal(err.Error())
-	}
-
-	return nil
-}

newconfig/config.go

@@ -1,258 +0,0 @@
-package newconfig
-
-import (
-	"bytes"
-	"encoding/json"
-	"net"
-	"strconv"
-	"time"
-
-	"github.com/juju/errors"
-	"go.uber.org/zap"
-	statsd "gopkg.in/alexcesaro/statsd.v2"
-)
-
-type SecretMode uint8
-
-func (s SecretMode) String() string {
-	switch s {
-	case SecretModeSimple:
-		return "simple"
-	case SecretModeSecured:
-		return "secured"
-	}
-	return "tls"
-}
-
-const (
-	SecretModeSimple SecretMode = iota
-	SecretModeSecured
-	SecretModeTLS
-)
-
-const SimpleSecretLength = 16
-
-type OptionType uint8
-
-const (
-	OptionTypeDebug OptionType = iota
-	OptionTypeVerbose
-
-	OptionTypeBindIP
-	OptionTypeBindPort
-	OptionTypePublicIPv4
-	OptionTypePublicIPv4Port
-	OptionTypePublicIPv6
-	OptionTypePublicIPv6Port
-	OptionTypeStatsIP
-	OptionTypeStatsPort
-
-	OptionTypeStatsdIP
-	OptionTypeStatsdPort
-	OptionTypeStatsdNetwork
-	OptionTypeStatsdPrefix
-	OptionTypeStatsdTagsFormat
-	OptionTypeStatsdTags
-	OptionTypePrometheusPrefix
-
-	OptionTypeWriteBufferSize
-	OptionTypeReadBufferSize
-
-	OptionTypeAntiReplayMaxSize
-	OptionTypeAntiReplayEvictionTime
-
-	OptionTypeSecret
-	OptionTypeAdtag
-)
-
-type BufferSize struct {
-	Read  int `json:"read"`
-	Write int `json:"write"`
-}
-
-type AntiReplay struct {
-	MaxSize      int           `json:"max_size"`
-	EvictionTime time.Duration `json:"duration"`
-}
-
-type Stats struct {
-	Prefix string `json:"prefix"`
-}
-
-type StatsdStats struct {
-	Stats
-
-	Addr       Addr              `json:"addr"`
-	Tags       map[string]string `json:"tags"`
-	TagsFormat statsd.TagFormat  `json:"format"`
-}
-
-type PrometheusStats struct {
-	Stats
-}
-
-type Addr struct {
-	IP   net.IP `json:"ip"`
-	Port int    `json:"port"`
-	net  string
-}
-
-func (a Addr) Network() string {
-	if a.net == "" {
-		return "tcp"
-	}
-	return a.net
-}
-
-func (a Addr) String() string {
-	return net.JoinHostPort(a.IP.String(), strconv.Itoa(a.Port))
-}
-
-func (a Addr) MarshalJSON() ([]byte, error) {
-	data := map[string]string{
-		"network": a.Network(),
-		"addr":    a.String(),
-	}
-	return json.Marshal(data)
-}
-
-type Config struct {
-	BufferSize BufferSize `json:"buffer_size"`
-	AntiReplay AntiReplay `json:"anti_replay"`
-
-	ListenAddr     Addr `json:"listen_addr"`
-	PublicIPv4Addr Addr `json:"public_ipv4_addr"`
-	PublicIPv6Addr Addr `json:"public_ipv6_addr"`
-	StatsAddr      Addr `json:"stats_addr"`
-
-	StatsdStats     StatsdStats     `json:"stats_statsd"`
-	PrometheusStats PrometheusStats `json:"stats_prometheus"`
-
-	Debug      bool       `json:"debug"`
-	Verbose    bool       `json:"verbose"`
-	SecretMode SecretMode `json:"secret_mode"`
-	Secret     []byte     `json:"secret"`
-	AdTag      []byte     `json:"adtag"`
-}
-
-func (c Config) String() string {
-	data, _ := json.Marshal(c)
-	return string(data)
-}
-
-type ConfigOpt struct {
-	Option OptionType
-	Value  interface{}
-}
-
-var C = Config{}
-
-func Init(options ...ConfigOpt) error { // nolint: gocyclo
-	for _, opt := range options {
-		switch opt.Option {
-		case OptionTypeDebug:
-			C.Debug = opt.Value.(bool)
-		case OptionTypeVerbose:
-			C.Verbose = opt.Value.(bool)
-		case OptionTypeBindIP:
-			C.ListenAddr.IP = opt.Value.(net.IP)
-		case OptionTypeBindPort:
-			C.ListenAddr.Port = int(opt.Value.(uint16))
-		case OptionTypePublicIPv4:
-			C.PublicIPv4Addr.IP = opt.Value.(net.IP)
-		case OptionTypePublicIPv4Port:
-			C.PublicIPv4Addr.Port = int(opt.Value.(uint16))
-		case OptionTypePublicIPv6:
-			C.PublicIPv6Addr.IP = opt.Value.(net.IP)
-		case OptionTypePublicIPv6Port:
-			C.PublicIPv6Addr.Port = int(opt.Value.(uint16))
-		case OptionTypeStatsIP:
-			C.StatsAddr.IP = opt.Value.(net.IP)
-		case OptionTypeStatsPort:
-			C.StatsAddr.Port = int(opt.Value.(uint16))
-		case OptionTypeStatsdIP:
-			C.StatsdStats.Addr.IP = opt.Value.(net.IP)
-		case OptionTypeStatsdPort:
-			C.StatsdStats.Addr.Port = int(opt.Value.(uint16))
-		case OptionTypeStatsdNetwork:
-			C.StatsdStats.Addr.net = opt.Value.(string)
-		case OptionTypeStatsdPrefix:
-			C.StatsdStats.Prefix = opt.Value.(string)
-		case OptionTypeStatsdTagsFormat:
-			value := opt.Value.(string)
-			switch value {
-			case "datadog":
-				C.StatsdStats.TagsFormat = statsd.Datadog
-			case "influxdb":
-				C.StatsdStats.TagsFormat = statsd.InfluxDB
-			default:
-				return errors.Errorf("Incorrect statsd tag %s", value)
-			}
-		case OptionTypeStatsdTags:
-			C.StatsdStats.Tags = opt.Value.(map[string]string)
-		case OptionTypePrometheusPrefix:
-			C.PrometheusStats.Prefix = opt.Value.(string)
-		case OptionTypeWriteBufferSize:
-			C.BufferSize.Write = int(opt.Value.(uint32))
-		case OptionTypeReadBufferSize:
-			C.BufferSize.Read = int(opt.Value.(uint32))
-		case OptionTypeAntiReplayMaxSize:
-			C.AntiReplay.MaxSize = opt.Value.(int)
-		case OptionTypeAntiReplayEvictionTime:
-			C.AntiReplay.EvictionTime = opt.Value.(time.Duration)
-		case OptionTypeSecret:
-			C.Secret = opt.Value.([]byte)
-		case OptionTypeAdtag:
-			C.AdTag = opt.Value.([]byte)
-		default:
-			return errors.Errorf("Unknown tag %v", opt.Option)
-		}
-	}
-
-	switch {
-	case len(C.Secret) == 1+SimpleSecretLength && bytes.HasPrefix(C.Secret, []byte{0xdd}):
-		C.SecretMode = SecretModeSecured
-		C.Secret = bytes.TrimPrefix(C.Secret, []byte{0xdd})
-	case len(C.Secret) == SimpleSecretLength:
-		C.SecretMode = SecretModeSimple
-	default:
-		return errors.New("Incorrect secret")
-	}
-
-	return nil
-}
-
-func InitPublicAddress() error {
-	if C.PublicIPv4Addr.Port == 0 {
-		C.PublicIPv4Addr.Port = C.ListenAddr.Port
-	}
-	if C.PublicIPv6Addr.Port == 0 {
-		C.PublicIPv6Addr.Port = C.ListenAddr.Port
-	}
-
-	foundAddress := C.PublicIPv4Addr.IP != nil || C.PublicIPv6Addr.IP != nil
-	if C.PublicIPv4Addr.IP == nil {
-		ip, err := getGlobalIPv4()
-		if err != nil {
-			zap.S().Warnw("Cannot resolve public address", "error", err)
-		} else {
-			C.PublicIPv4Addr.IP = ip
-			foundAddress = true
-		}
-	}
-	if C.PublicIPv6Addr.IP == nil {
-		ip, err := getGlobalIPv6()
-		if err != nil {
-			zap.S().Warnw("Cannot resolve public address", "error", err)
-		} else {
-			C.PublicIPv6Addr.IP = ip
-			foundAddress = true
-		}
-	}
-
-	if !foundAddress {
-		return errors.New("Cannot resolve any public address")
-	}
-
-	return nil
-}

newobfuscated2/client_protocol.go

@@ -1,95 +0,0 @@
-package newobfuscated2
-
-import (
-	"bytes"
-	"crypto/aes"
-	"crypto/cipher"
-	"crypto/sha256"
-	"encoding/binary"
-	"io"
-	"time"
-
-	"github.com/juju/errors"
-
-	"github.com/9seconds/mtg/newantireplay"
-	"github.com/9seconds/mtg/newconfig"
-	"github.com/9seconds/mtg/newprotocol"
-	"github.com/9seconds/mtg/newwrappers"
-)
-
-const clientProtocolHandshakeTimeout = 10 * time.Second
-
-type ClientProtocol struct {
-	newprotocol.BaseProtocol
-}
-
-func (c *ClientProtocol) Handshake(socket newwrappers.StreamReadWriteCloser) (newwrappers.StreamReadWriteCloser, error) {
-	fm, err := c.ReadFrame(socket)
-	if err != nil {
-		return nil, errors.Annotate(err, "Cannot make client handshake")
-	}
-
-	decHasher := sha256.New()
-	decHasher.Write(fm.key())           // nolint: errcheck
-	decHasher.Write(newconfig.C.Secret) // nolint: errcheck
-	decryptor := makeStreamCipher(decHasher.Sum(nil), fm.iv())
-
-	invertedFrame := fm.invert()
-	encHasher := sha256.New()
-	encHasher.Write(invertedFrame.key()) // nolint: errcheck
-	encHasher.Write(newconfig.C.Secret)  // nolint: errcheck
-	encryptor := makeStreamCipher(encHasher.Sum(nil), invertedFrame.iv())
-
-	decryptedFrame := frame{}
-	decryptor.XORKeyStream(decryptedFrame.bytes(), fm.bytes())
-
-	magic := decryptedFrame.magic()
-	switch {
-	case bytes.Equal(magic, newprotocol.ConnectionTagAbridged):
-		c.ConnectionType = newprotocol.ConnectionTypeAbridged
-	case bytes.Equal(magic, newprotocol.ConnectionTagIntermediate):
-		c.ConnectionType = newprotocol.ConnectionTypeIntermediate
-	case bytes.Equal(magic, newprotocol.ConnectionTagSecure):
-		c.ConnectionType = newprotocol.ConnectionTypeSecure
-	default:
-		return nil, errors.New("Unknown connection type")
-	}
-
-	c.ConnectionProtocol = newprotocol.ConnectionProtocolIPv4
-	if socket.LocalAddr().IP.To4() == nil {
-		c.ConnectionProtocol = newprotocol.ConnectionProtocolIPv6
-	}
-
-	buf := bytes.NewReader(decryptedFrame.dc())
-	if err := binary.Read(buf, binary.LittleEndian, &c.DC); err != nil {
-		c.DC = 1
-	}
-
-	antiReplayKey := decryptedFrame.unique()
-	if newantireplay.Has(antiReplayKey) {
-		return nil, errors.New("Replay attack is detected")
-	}
-	newantireplay.Add(antiReplayKey)
-
-	return newwrappers.NewObfuscated2(socket, encryptor, decryptor), nil
-}
-
-func (c *ClientProtocol) ReadFrame(socket newwrappers.StreamReader) (fm frame, err error) {
-	if _, err := io.ReadFull(handshakeReader{socket}, fm.bytes()); err != nil {
-		err = errors.Annotate(err, "Cannot extract obfuscated2 frame")
-	}
-	return
-}
-
-type handshakeReader struct {
-	parent newwrappers.StreamReader
-}
-
-func (h handshakeReader) Read(p []byte) (int, error) {
-	return h.parent.ReadTimeout(p, clientProtocolHandshakeTimeout)
-}
-
-func makeStreamCipher(key, iv []byte) cipher.Stream {
-	block, _ := aes.NewCipher(key) // nolint: gosec
-	return cipher.NewCTR(block, iv)
-}

newobfuscated2/frame.go

@@ -1,54 +0,0 @@
-package newobfuscated2
-
-const (
-	frameLenKey   = 32
-	frameLenIV    = 16
-	frameLenMagic = 4
-	frameLenDC    = 2
-
-	frameOffsetFirst = 8
-	frameOffsetKey   = frameOffsetFirst + frameLenKey
-	frameOffsetIV    = frameOffsetKey + frameLenIV
-	frameOffsetMagic = frameOffsetIV + frameLenMagic
-	frameOffsetDC    = frameOffsetMagic + frameLenDC
-
-	frameLen = 64
-)
-
-// [frameOffsetFirst:frameOffsetKey:frameOffsetIV:frameOffsetMagic:frameOffsetDC:frameOffsetEnd]
-type frame struct {
-	data [frameLen]byte
-}
-
-func (f *frame) bytes() []byte {
-	return f.data[:]
-}
-
-func (f *frame) key() []byte {
-	return f.data[frameOffsetFirst:frameOffsetKey]
-}
-
-func (f *frame) iv() []byte {
-	return f.data[frameOffsetKey:frameOffsetIV]
-}
-
-func (f *frame) magic() []byte {
-	return f.data[frameOffsetIV:frameOffsetMagic]
-}
-
-func (f *frame) dc() []byte {
-	return f.data[frameOffsetMagic:frameOffsetDC]
-}
-
-func (f *frame) unique() []byte {
-	return f.data[frameOffsetFirst:frameOffsetDC]
-}
-
-func (f *frame) invert() (nf frame) {
-	nf = *f
-	for i := 0; i < frameLenKey+frameLenIV; i++ {
-		nf.data[frameOffsetFirst+i] = nf.data[frameOffsetIV-1-i]
-	}
-
-	return
-}

newobfuscated2/telegram_protocol.go

@@ -1,61 +0,0 @@
-package newobfuscated2
-
-import (
-	"crypto/rand"
-
-	"github.com/juju/errors"
-
-	"github.com/9seconds/mtg/newprotocol"
-	"github.com/9seconds/mtg/newwrappers"
-)
-
-type TelegramProtocol struct {
-	newprotocol.BaseProtocol
-}
-
-func (t *TelegramProtocol) Handshake(socketRaw newwrappers.Wrap, client *ClientProtocol) (newwrappers.StreamReadWriteCloser, error) {
-	socket := socketRaw.(newwrappers.StreamReadWriteCloser)
-	fm := generateFrame(client)
-	data := fm.bytes()
-
-	encryptor := makeStreamCipher(fm.key(), fm.iv())
-	decryptedFrame := fm.invert()
-	decryptor := makeStreamCipher(decryptedFrame.key(), decryptedFrame.iv())
-
-	copyFrame := make([]byte, frameLen)
-	copy(copyFrame[:frameOffsetIV], data[:frameOffsetIV])
-	encryptor.XORKeyStream(data, data)
-	copy(data[:frameOffsetIV], copyFrame[:frameOffsetIV])
-
-	if _, err := socket.Write(data); err != nil {
-		return nil, errors.Annotate(err, "Cannot write handshate frame to Telegram")
-	}
-
-	return newwrappers.NewObfuscated2(socket, encryptor, decryptor), nil
-}
-
-func generateFrame(client *ClientProtocol) (fm frame) {
-	for {
-		data := fm.bytes()
-		if _, err := rand.Read(data); err != nil {
-			continue
-		}
-		if data[0] == 0xef {
-			continue
-		}
-
-		val := (uint32(data[3]) << 24) | (uint32(data[2]) << 16) | (uint32(data[1]) << 8) | uint32(data[0])
-		if val == 0x44414548 || val == 0x54534f50 || val == 0x20544547 || val == 0x4954504f || val == 0xeeeeeeee {
-			continue
-		}
-
-		val = (uint32(data[7]) << 24) | (uint32(data[6]) << 16) | (uint32(data[5]) << 8) | uint32(data[4])
-		if val == 0x00000000 {
-			continue
-		}
-
-		copy(fm.magic(), client.ConnectionType.Tag())
-
-		return
-	}
-}

newprotocol/base_protocol.go

@@ -1,7 +0,0 @@
-package newprotocol
-
-type BaseProtocol struct {
-	ConnectionType     ConnectionType
-	ConnectionProtocol ConnectionProtocol
-	DC                 int16
-}

newproxy/proxy.go

@@ -1 +0,0 @@
-package newproxy

newstats/stats.go

@@ -1,93 +0,0 @@
-package newstats
-
-import (
-	"net"
-	"net/http"
-
-	"github.com/juju/errors"
-
-	"github.com/9seconds/mtg/newconfig"
-	"github.com/9seconds/mtg/newprotocol"
-)
-
-type Stats interface {
-	IngressTraffic(int)
-	EgressTraffic(int)
-	ClientConnected(newprotocol.ConnectionType, *net.TCPAddr)
-	ClientDisconnected(newprotocol.ConnectionType, *net.TCPAddr)
-	Crash()
-	AntiReplayDetected()
-}
-
-type multiStats []Stats
-
-func (m multiStats) IngressTraffic(traffic int) {
-	for i := range m {
-		go m[i].IngressTraffic(traffic)
-	}
-}
-
-func (m multiStats) EgressTraffic(traffic int) {
-	for i := range m {
-		go m[i].EgressTraffic(traffic)
-	}
-}
-
-func (m multiStats) ClientConnected(connectionType newprotocol.ConnectionType, addr *net.TCPAddr) {
-	for i := range m {
-		go m[i].ClientConnected(connectionType, addr)
-	}
-}
-
-func (m multiStats) ClientDisconnected(connectionType newprotocol.ConnectionType, addr *net.TCPAddr) {
-	for i := range m {
-		go m[i].ClientDisconnected(connectionType, addr)
-	}
-}
-
-func (m multiStats) Crash() {
-	for i := range m {
-		go m[i].Crash()
-	}
-}
-
-func (m multiStats) AntiReplayDetected() {
-	for i := range m {
-		go m[i].AntiReplayDetected()
-	}
-}
-
-var S Stats
-
-func Init() error {
-	mux := http.NewServeMux()
-
-	instanceJSON := newStatsJSON(mux)
-	instancePrometheus, err := newStatsPrometheus(mux)
-	if err != nil {
-		return errors.Annotate(err, "Cannot initialize Prometheus")
-	}
-
-	stats := []Stats{instanceJSON, instancePrometheus}
-	if newconfig.C.StatsdStats.Addr.IP != nil {
-		instanceStatsd, err := newStatsStatsd()
-		if err != nil {
-			return errors.Annotate(err, "Cannot initialize StatsD")
-		}
-		stats = append(stats, instanceStatsd)
-	}
-
-	listener, err := net.Listen("tcp", newconfig.C.StatsAddr.String())
-	if err != nil {
-		return errors.Annotate(err, "Cannot initialize stats server")
-	}
-
-	srv := http.Server{
-		Handler: mux,
-	}
-	go srv.Serve(listener) // nolint: errcheck
-
-	S = multiStats(stats)
-
-	return nil
-}

obfuscated2/client_protocol.go

@@ -0,0 +1,94 @@
+package obfuscated2
+
+import (
+	"bytes"
+	"crypto/sha256"
+	"encoding/binary"
+	"io"
+	"time"
+
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/antireplay"
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/protocol"
+	"github.com/9seconds/mtg/utils"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+const clientProtocolHandshakeTimeout = 10 * time.Second
+
+type ClientProtocol struct {
+	protocol.BaseProtocol
+}
+
+func (c *ClientProtocol) Handshake(socket wrappers.StreamReadWriteCloser) (wrappers.StreamReadWriteCloser, error) {
+	fm, err := c.ReadFrame(socket)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot make client handshake")
+	}
+
+	decHasher := sha256.New()
+	decHasher.Write(fm.Key())        // nolint: errcheck
+	decHasher.Write(config.C.Secret) // nolint: errcheck
+	decryptor := utils.MakeStreamCipher(decHasher.Sum(nil), fm.IV())
+
+	invertedFrame := fm.Invert()
+	encHasher := sha256.New()
+	encHasher.Write(invertedFrame.Key()) // nolint: errcheck
+	encHasher.Write(config.C.Secret)     // nolint: errcheck
+	encryptor := utils.MakeStreamCipher(encHasher.Sum(nil), invertedFrame.IV())
+
+	decryptedFrame := Frame{}
+	decryptor.XORKeyStream(decryptedFrame.Bytes(), fm.Bytes())
+
+	magic := decryptedFrame.Magic()
+	switch {
+	case bytes.Equal(magic, conntypes.ConnectionTagAbridged):
+		c.ConnectionType = conntypes.ConnectionTypeAbridged
+	case bytes.Equal(magic, conntypes.ConnectionTagIntermediate):
+		c.ConnectionType = conntypes.ConnectionTypeIntermediate
+	case bytes.Equal(magic, conntypes.ConnectionTagSecure):
+		c.ConnectionType = conntypes.ConnectionTypeSecure
+	default:
+		return nil, errors.New("Unknown connection type")
+	}
+
+	c.ConnectionProtocol = conntypes.ConnectionProtocolIPv4
+	if socket.LocalAddr().IP.To4() == nil {
+		c.ConnectionProtocol = conntypes.ConnectionProtocolIPv6
+	}
+
+	buf := bytes.NewReader(decryptedFrame.DC())
+	if err := binary.Read(buf, binary.LittleEndian, &c.DC); err != nil {
+		c.DC = conntypes.DCDefaultIdx
+	}
+
+	antiReplayKey := decryptedFrame.Unique()
+	if antireplay.Has(antiReplayKey) {
+		return nil, errors.New("Replay attack is detected")
+	}
+	antireplay.Add(antiReplayKey)
+
+	return wrappers.NewObfuscated2(socket, encryptor, decryptor), nil
+}
+
+func (c *ClientProtocol) ReadFrame(socket wrappers.StreamReader) (fm Frame, err error) {
+	if _, err = io.ReadFull(handshakeReader{socket}, fm.Bytes()); err != nil {
+		err = errors.Annotate(err, "Cannot extract obfuscated2 frame")
+	}
+	return
+}
+
+type handshakeReader struct {
+	parent wrappers.StreamReader
+}
+
+func (h handshakeReader) Read(p []byte) (int, error) {
+	return h.parent.ReadTimeout(p, clientProtocolHandshakeTimeout)
+}
+
+func MakeClientProtocol() protocol.ClientProtocol {
+	return &ClientProtocol{}
+}

obfuscated2/frame.go

@@ -1,17 +1,5 @@
 package obfuscated2
 
-import (
-	"bytes"
-	"crypto/rand"
-	"encoding/binary"
-	"io"
-
-	"github.com/juju/errors"
-
-	"github.com/9seconds/mtg/mtproto"
-)
-
-// [frameOffsetFirst:frameOffsetKey:frameOffsetIV:frameOffsetMagic:frameOffsetDC:frameOffsetEnd]
 const (
 	frameLenKey   = 32
 	frameLenIV    = 16
@@ -24,98 +12,43 @@ const (
 	frameOffsetMagic = frameOffsetIV + frameLenMagic
 	frameOffsetDC    = frameOffsetMagic + frameLenDC
 
-	FrameLen = 64
+	frameLen = 64
 )
 
-// Frame represents handshake frame. Telegram sends 64 bytes of obfuscated2
-// initialization data first.
-// https://blog.susanka.eu/how-telegram-obfuscates-its-mtproto-traffic/
-type Frame []byte
-
-// Key returns AES encryption key.
-func (f Frame) Key() []byte {
-	return f[frameOffsetFirst:frameOffsetKey]
+// [frameOffsetFirst:frameOffsetKey:frameOffsetIV:frameOffsetMagic:frameOffsetDC:frameOffsetEnd]
+type Frame struct {
+	data [frameLen]byte
 }
 
-// IV returns AES encryption initialization vector
-func (f Frame) IV() []byte {
-	return f[frameOffsetKey:frameOffsetIV]
+func (f *Frame) Bytes() []byte {
+	return f.data[:]
 }
 
-// Magic returns magic bytes from last 8 bytes of frame. Telegram checks
-// for values there. If after decryption magic is not as expected,
-// connection considered as failed.
-func (f Frame) Magic() []byte {
-	return f[frameOffsetIV:frameOffsetMagic]
+func (f *Frame) Key() []byte {
+	return f.data[frameOffsetFirst:frameOffsetKey]
 }
 
-// DC returns number of datacenter IP client wants to use.
-func (f Frame) DC() (n int16) {
-	buf := bytes.NewReader(f[frameOffsetMagic:frameOffsetDC])
-	if err := binary.Read(buf, binary.LittleEndian, &n); err != nil {
-		n = 1
-	}
-
-	return
+func (f *Frame) IV() []byte {
+	return f.data[frameOffsetKey:frameOffsetIV]
 }
 
-// ConnectionType identifies connection type of the handshake frame.
-func (f Frame) ConnectionType() (mtproto.ConnectionType, error) {
-	return mtproto.ConnectionTagFromHandshake(f.Magic())
+func (f *Frame) Magic() []byte {
+	return f.data[frameOffsetIV:frameOffsetMagic]
 }
 
-// Invert inverts frame for extracting encryption keys. Pkease check that link:
-// https://blog.susanka.eu/how-telegram-obfuscates-its-mtproto-traffic/
-func (f Frame) Invert() Frame {
-	reversed := make(Frame, FrameLen)
-	copy(reversed, f)
-
-	for i := 0; i < frameLenKey+frameLenIV; i++ {
-		reversed[frameOffsetFirst+i] = f[frameOffsetIV-1-i]
-	}
-
-	return reversed
+func (f *Frame) DC() []byte {
+	return f.data[frameOffsetMagic:frameOffsetDC]
 }
 
-// ExtractFrame extracts exact obfuscated2 handshake frame from given reader.
-func ExtractFrame(conn io.Reader) (Frame, error) {
-	frame := make(Frame, FrameLen)
-	buf := bytes.NewBuffer(frame)
-	buf.Reset()
-
-	if _, err := io.CopyN(buf, conn, FrameLen); err != nil {
-		return nil, errors.Annotate(err, "Cannot extract obfuscated header")
-	}
-	copy(frame, buf.Bytes())
-
-	return frame, nil
+func (f *Frame) Unique() []byte {
+	return f.data[frameOffsetFirst:frameOffsetDC]
 }
 
-func generateFrame(connectionType mtproto.ConnectionType) Frame {
-	frame := make(Frame, FrameLen)
-
-	for {
-		if _, err := rand.Read(frame); err != nil {
-			continue
-		}
-		if frame[0] == 0xef {
-			continue
-		}
-
-		val := (uint32(frame[3]) << 24) | (uint32(frame[2]) << 16) | (uint32(frame[1]) << 8) | uint32(frame[0])
-		if val == 0x44414548 || val == 0x54534f50 || val == 0x20544547 || val == 0x4954504f || val == 0xeeeeeeee {
-			continue
-		}
-
-		val = (uint32(frame[7]) << 24) | (uint32(frame[6]) << 16) | (uint32(frame[5]) << 8) | uint32(frame[4])
-		if val == 0x00000000 {
-			continue
-		}
-
-		// error has to be checked before calling this function
-		tag, _ := connectionType.Tag() // nolint: errcheck, gosec
-		copy(frame.Magic(), tag)
-
-		return frame
+func (f *Frame) Invert() (nf Frame) {
+	nf = *f
+	for i := 0; i < frameLenKey+frameLenIV; i++ {
+		nf.data[frameOffsetFirst+i] = f.data[frameOffsetIV-1-i]
 	}
+
+	return
 }

obfuscated2/telegram_protocol.go

@@ -0,0 +1,78 @@
+package obfuscated2
+
+import (
+	"crypto/rand"
+
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/protocol"
+	"github.com/9seconds/mtg/telegram"
+	"github.com/9seconds/mtg/utils"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+type TelegramProtocol struct {
+	protocol.BaseProtocol
+
+	dialer telegram.Telegram
+}
+
+func (t *TelegramProtocol) Handshake(req *protocol.TelegramRequest) (wrappers.Wrap, error) {
+	socket, err := t.dialer.Dial(req.Ctx,
+		req.Cancel,
+		req.ClientProtocol.GetDC(),
+		req.ClientProtocol.GetConnectionProtocol())
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot dial to Telegram")
+	}
+	fm := generateFrame(req.ClientProtocol)
+	data := fm.Bytes()
+
+	encryptor := utils.MakeStreamCipher(fm.Key(), fm.IV())
+	decryptedFrame := fm.Invert()
+	decryptor := utils.MakeStreamCipher(decryptedFrame.Key(), decryptedFrame.IV())
+
+	copyFrame := make([]byte, frameLen)
+	copy(copyFrame[:frameOffsetIV], data[:frameOffsetIV])
+	encryptor.XORKeyStream(data, data)
+	copy(data[:frameOffsetIV], copyFrame[:frameOffsetIV])
+
+	if _, err := socket.Write(data); err != nil {
+		return nil, errors.Annotate(err, "Cannot write handshate frame to Telegram")
+	}
+
+	return wrappers.NewObfuscated2(socket, encryptor, decryptor), nil
+}
+
+func MakeTelegramProtocol(dialer telegram.Telegram) protocol.TelegramProtocol {
+	return &TelegramProtocol{
+		dialer: dialer,
+	}
+}
+
+func generateFrame(cp protocol.ClientProtocol) (fm Frame) {
+	data := fm.Bytes()
+
+	for {
+		if _, err := rand.Read(data); err != nil {
+			continue
+		}
+		if data[0] == 0xef {
+			continue
+		}
+
+		val := (uint32(data[3]) << 24) | (uint32(data[2]) << 16) | (uint32(data[1]) << 8) | uint32(data[0])
+		if val == 0x44414548 || val == 0x54534f50 || val == 0x20544547 || val == 0x4954504f || val == 0xeeeeeeee {
+			continue
+		}
+
+		val = (uint32(data[7]) << 24) | (uint32(data[6]) << 16) | (uint32(data[5]) << 8) | uint32(data[4])
+		if val == 0x00000000 {
+			continue
+		}
+
+		copy(fm.Magic(), cp.GetConnectionType().Tag())
+
+		return
+	}
+}

protocol/base_protocol.go

@@ -0,0 +1,21 @@
+package protocol
+
+import "github.com/9seconds/mtg/conntypes"
+
+type BaseProtocol struct {
+	ConnectionType     conntypes.ConnectionType
+	ConnectionProtocol conntypes.ConnectionProtocol
+	DC                 conntypes.DC
+}
+
+func (b *BaseProtocol) GetConnectionType() conntypes.ConnectionType {
+	return b.ConnectionType
+}
+
+func (b *BaseProtocol) GetConnectionProtocol() conntypes.ConnectionProtocol {
+	return b.ConnectionProtocol
+}
+
+func (b *BaseProtocol) GetDC() conntypes.DC {
+	return b.DC
+}

protocol/interfaces.go

@@ -0,0 +1,22 @@
+package protocol
+
+import (
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/telegram"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+type ClientProtocol interface {
+	Handshake(wrappers.StreamReadWriteCloser) (wrappers.StreamReadWriteCloser, error)
+	GetConnectionType() conntypes.ConnectionType
+	GetConnectionProtocol() conntypes.ConnectionProtocol
+	GetDC() conntypes.DC
+}
+
+type ClientProtocolMaker func() ClientProtocol
+
+type TelegramProtocol interface {
+	Handshake(*TelegramRequest) (wrappers.Wrap, error)
+}
+
+type TelegramProtocolMaker func(telegram.Telegram) TelegramProtocol

protocol/request.go

@@ -0,0 +1,18 @@
+package protocol
+
+import (
+	"context"
+
+	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/wrappers"
+)
+
+type TelegramRequest struct {
+	Logger         *zap.SugaredLogger
+	ClientConn     wrappers.StreamReadWriteCloser
+	ConnID         wrappers.ConnID
+	Ctx            context.Context
+	Cancel         context.CancelFunc
+	ClientProtocol ClientProtocol
+}

proxy/proxy.go

@@ -6,173 +6,123 @@ import (
 	"net"
 	"sync"
 
-	"github.com/gofrs/uuid"
-	"github.com/juju/errors"
 	"go.uber.org/zap"
 
-	"github.com/9seconds/mtg/antireplay"
-	"github.com/9seconds/mtg/client"
 	"github.com/9seconds/mtg/config"
-	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/protocol"
 	"github.com/9seconds/mtg/stats"
 	"github.com/9seconds/mtg/telegram"
+	"github.com/9seconds/mtg/utils"
 	"github.com/9seconds/mtg/wrappers"
 )
 
-// Proxy is a core of this program.
+const directPipeBufferSize = 1024 * 1024
+
 type Proxy struct {
-	antiReplayCache antireplay.Cache
-	clientInit      client.Init
-	tg              telegram.Telegram
-	conf            *config.Config
+	Logger                *zap.SugaredLogger
+	ClientProtocolMaker   protocol.ClientProtocolMaker
+	TelegramProtocolMaker protocol.TelegramProtocolMaker
+	TelegramDialer        telegram.Telegram
 }
 
-// Serve runs TCP proxy server.
-func (p *Proxy) Serve() error {
-	lsock, err := net.Listen("tcp", p.conf.BindAddr())
-	if err != nil {
-		return errors.Annotate(err, "Cannot create listen socket")
-	}
-
+func (p *Proxy) Serve(listener net.Listener) {
 	for {
-		if conn, err := lsock.Accept(); err != nil {
-			zap.S().Errorw("Cannot allocate incoming connection", "error", err)
-		} else {
-			go p.accept(conn)
+		conn, err := listener.Accept()
+		if err != nil {
+			p.Logger.Errorw("Cannot allocate incoming connection", "error", err)
+			continue
 		}
+		go p.accept(conn)
 	}
 }
 
 func (p *Proxy) accept(conn net.Conn) {
-	connID := uuid.Must(uuid.NewV4()).String()
-	log := zap.S().With("connection_id", connID).Named("main")
-	ctx, cancel := context.WithCancel(context.Background())
-
 	defer func() {
-		cancel()
-		conn.Close() // nolint: errcheck, gosec
-
+		conn.Close()
 		if err := recover(); err != nil {
-			stats.NewCrash()
-			log.Errorw("Crash of accept handler", "error", err)
+			stats.S.Crash()
+			p.Logger.Errorw("Crash of accept handler", "error", err)
 		}
 	}()
 
-	log.Infow("Client connected", "addr", conn.RemoteAddr())
+	connID := wrappers.NewConnID()
+	logger := p.Logger.With("connection_id", connID)
 
-	clientConn, opts, err := p.clientInit(ctx, cancel, conn, connID, p.antiReplayCache, p.conf)
-	if err != nil {
-		log.Errorw("Cannot initialize client connection", "error", err)
+	if err := utils.InitTCP(conn); err != nil {
+		logger.Errorw("Cannot initialize client TCP connection", "error", err)
 		return
 	}
-	defer clientConn.(io.Closer).Close() // nolint: errcheck
 
-	if p.conf.SecureOnly && opts.ConnectionType != mtproto.ConnectionTypeSecure {
-		log.Errorw("Proxy supports only secure connections", "connection_type", opts.ConnectionType)
-		return
-	}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
-	stats.ClientConnected(opts.ConnectionType, clientConn.RemoteAddr())
-	defer stats.ClientDisconnected(opts.ConnectionType, clientConn.RemoteAddr())
+	wrappedConn := wrappers.NewClientConn(ctx, cancel, conn, connID)
+	wrappedConn = wrappers.NewTraffic(wrappedConn)
+	defer wrappedConn.Close()
 
-	serverConn, err := p.getTelegramConn(ctx, cancel, opts, connID)
+	clientProtocol := p.ClientProtocolMaker()
+	wrappedConn, err := clientProtocol.Handshake(wrappedConn)
 	if err != nil {
-		log.Errorw("Cannot initialize server connection", "error", err)
+		logger.Warnw("Cannot perform client handshake", "error", err)
 		return
 	}
-	defer serverConn.(io.Closer).Close() // nolint: errcheck
-
-	go func() {
-		<-ctx.Done()
-		serverConn.(io.Closer).Close() // nolint: gosec
-		clientConn.(io.Closer).Close() // nolint: gosec
-	}()
-
-	wait := &sync.WaitGroup{}
-	wait.Add(2)
+	defer wrappedConn.Close()
+
+	stats.S.ClientConnected(clientProtocol.GetConnectionType(), wrappedConn.RemoteAddr())
+	defer stats.S.ClientDisconnected(clientProtocol.GetConnectionType(), wrappedConn.RemoteAddr())
+	logger.Infow("Client connected", "addr", conn.RemoteAddr())
+
+	req := &protocol.TelegramRequest{
+		Logger:         logger,
+		ClientConn:     wrappedConn,
+		ConnID:         connID,
+		Ctx:            ctx,
+		Cancel:         cancel,
+		ClientProtocol: clientProtocol,
+	}
 
-	if p.conf.UseMiddleProxy() {
-		clientPacket := clientConn.(wrappers.PacketReadWriteCloser)
-		serverPacket := serverConn.(wrappers.PacketReadWriteCloser)
-		go p.middlePipe(clientPacket, serverPacket, wait, &opts.ReadHacks)
-		p.middlePipe(serverPacket, clientPacket, wait, &opts.WriteHacks)
+	if len(config.C.AdTag) > 0 {
+		err = p.acceptMiddleProxyConnection(req)
 	} else {
-		clientStream := clientConn.(wrappers.StreamReadWriteCloser)
-		serverStream := serverConn.(wrappers.StreamReadWriteCloser)
-		go p.directPipe(clientStream, serverStream, wait, p.conf.ReadBufferSize)
-		p.directPipe(serverStream, clientStream, wait, p.conf.WriteBufferSize)
+		err = p.acceptDirectConnection(req)
 	}
 
-	wait.Wait()
-
-	log.Infow("Client disconnected", "addr", conn.RemoteAddr())
+	logger.Infow("Client disconnected", "error", err, "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(ctx, cancel, connID, opts)
+func (p *Proxy) acceptDirectConnection(request *protocol.TelegramRequest) error {
+	telegramProtocol := p.TelegramProtocolMaker(p.TelegramDialer)
+	telegramConnRaw, err := telegramProtocol.Handshake(request)
 	if err != nil {
-		return nil, errors.Annotate(err, "Cannot dial to Telegram")
+		return err
 	}
+	telegramConn := telegramConnRaw.(wrappers.StreamReadWriteCloser)
+	defer telegramConn.Close()
 
-	packetConn, err := p.tg.Init(opts, streamConn)
-	if err != nil {
-		return nil, errors.Annotate(err, "Cannot handshake telegram")
-	}
-
-	return packetConn, nil
-}
+	wg := &sync.WaitGroup{}
+	wg.Add(2)
 
-func (p *Proxy) middlePipe(src wrappers.PacketReadCloser, dst io.Writer, wait *sync.WaitGroup, hacks *mtproto.Hacks) {
-	defer wait.Done()
+	go p.directPipe(telegramConn, request.ClientConn, wg, request.Logger)
+	go p.directPipe(request.ClientConn, telegramConn, wg, request.Logger)
 
-	for {
-		hacks.SimpleAck = false
-		hacks.QuickAck = false
+	<-request.Ctx.Done()
+	wg.Wait()
 
-		packet, err := src.Read()
-		if err != nil {
-			src.Logger().Warnw("Cannot read packet", "error", err)
-			return
-		}
-		if _, err = dst.Write(packet); err != nil {
-			src.Logger().Warnw("Cannot write packet", "error", err)
-			return
-		}
-	}
+	return request.Ctx.Err()
 }
 
-func (p *Proxy) directPipe(src wrappers.StreamReadCloser, dst io.Writer, wait *sync.WaitGroup, bufferSize int) {
-	defer wait.Done()
+func (p *Proxy) directPipe(dst io.Writer,
+	src io.Reader,
+	wg *sync.WaitGroup,
+	logger *zap.SugaredLogger) {
+	defer wg.Done()
 
-	buffer := make([]byte, bufferSize)
-	if _, err := io.CopyBuffer(dst, src, buffer); err != nil {
-		src.Logger().Warnw("Cannot pump sockets", "error", err)
+	buf := make([]byte, directPipeBufferSize)
+	if _, err := io.CopyBuffer(dst, src, buf); err != nil {
+		logger.Debugw("Cannot pump sockets", "error", err)
 	}
 }
 
-// NewProxy returns new proxy instance.
-func NewProxy(conf *config.Config) (*Proxy, error) {
-	var clientInit client.Init
-	var tg telegram.Telegram
-
-	cache, err := antireplay.NewCache(conf)
-	if err != nil {
-		return nil, errors.Annotate(err, "Cannot make proxy")
-	}
-
-	if conf.UseMiddleProxy() {
-		clientInit = client.MiddleInit
-		tg = telegram.NewMiddleTelegram(conf)
-	} else {
-		clientInit = client.DirectInit
-		tg = telegram.NewDirectTelegram(conf)
-	}
-
-	return &Proxy{
-		antiReplayCache: cache,
-		conf:            conf,
-		clientInit:      clientInit,
-		tg:              tg,
-	}, nil
+func (p *Proxy) acceptMiddleProxyConnection(request *protocol.TelegramRequest) error {
+	return nil
 }

stats/stats.go

@@ -1,175 +1,93 @@
 package stats
 
 import (
-	"encoding/json"
-	"fmt"
-	"strconv"
-	"time"
+	"net"
+	"net/http"
 
-	humanize "github.com/dustin/go-humanize"
+	"github.com/juju/errors"
 
 	"github.com/9seconds/mtg/config"
-	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/conntypes"
 )
 
-type uptime time.Time
-
-func (u uptime) MarshalJSON() ([]byte, error) {
-	duration := time.Since(time.Time(u))
-	value := map[string]string{
-		"seconds": strconv.Itoa(int(duration.Seconds())),
-		"human":   humanize.Time(time.Time(u)),
-	}
-
-	return json.Marshal(value)
-}
-
-type connectionType struct {
-	IPv6 uint32 `json:"ipv6"`
-	IPv4 uint32 `json:"ipv4"`
-}
-
-type baseConnections struct {
-	All          connectionType `json:"all"`
-	Abridged     connectionType `json:"abridged"`
-	Intermediate connectionType `json:"intermediate"`
-	Secure       connectionType `json:"secure"`
+type Stats interface {
+	IngressTraffic(int)
+	EgressTraffic(int)
+	ClientConnected(conntypes.ConnectionType, *net.TCPAddr)
+	ClientDisconnected(conntypes.ConnectionType, *net.TCPAddr)
+	Crash()
+	AntiReplayDetected()
 }
 
-type connections struct {
-	baseConnections
-}
+type multiStats []Stats
 
-func (c connections) MarshalJSON() ([]byte, error) {
-	c.All.IPv4 = c.Abridged.IPv4 + c.Intermediate.IPv4 + c.Secure.IPv4
-	c.All.IPv6 = c.Abridged.IPv6 + c.Intermediate.IPv6 + c.Secure.IPv6
-
-	return json.Marshal(c.baseConnections)
-}
-
-type traffic struct {
-	ingress uint64
-	egress  uint64
-}
-
-func (t *traffic) dumpValue(value uint64) map[string]interface{} {
-	return map[string]interface{}{
-		"bytes": value,
-		"human": humanize.Bytes(value),
+func (m multiStats) IngressTraffic(traffic int) {
+	for i := range m {
+		go m[i].IngressTraffic(traffic)
 	}
 }
 
-func (t traffic) MarshalJSON() ([]byte, error) {
-	value := map[string]map[string]interface{}{
-		"ingress": t.dumpValue(t.ingress),
-		"egress":  t.dumpValue(t.egress),
+func (m multiStats) EgressTraffic(traffic int) {
+	for i := range m {
+		go m[i].EgressTraffic(traffic)
 	}
-
-	return json.Marshal(value)
 }
 
-type speed struct {
-	ingress uint64
-	egress  uint64
-}
-
-func (s *speed) dumpValue(value uint64) map[string]interface{} {
-	return map[string]interface{}{
-		"bytes/s": value,
-		"human":   fmt.Sprintf("%s/s", humanize.Bytes(value)),
+func (m multiStats) ClientConnected(connectionType conntypes.ConnectionType, addr *net.TCPAddr) {
+	for i := range m {
+		go m[i].ClientConnected(connectionType, addr)
 	}
 }
 
-func (s speed) MarshalJSON() ([]byte, error) {
-	value := map[string]map[string]interface{}{
-		"ingress": s.dumpValue(s.ingress),
-		"egress":  s.dumpValue(s.egress),
+func (m multiStats) ClientDisconnected(connectionType conntypes.ConnectionType, addr *net.TCPAddr) {
+	for i := range m {
+		go m[i].ClientDisconnected(connectionType, addr)
 	}
-
-	return json.Marshal(value)
 }
 
-// Stats represents a statistics of the proxy.
-type Stats struct {
-	URLs        config.IPURLs `json:"urls"`
-	Connections connections   `json:"connections"`
-	Traffic     traffic       `json:"traffic"`
-	Speed       speed         `json:"speed"`
-	Uptime      uptime        `json:"uptime"`
-	Crashes     uint32        `json:"crashes"`
-
-	previousTraffic traffic
-}
-
-func (s *Stats) start() {
-	speedChan := time.Tick(time.Second)
-
-	for {
-		select {
-		case <-speedChan:
-			s.handleSpeed()
-		case event := <-trafficChan:
-			s.handleTraffic(event)
-		case event := <-connectionsChan:
-			s.handleConnection(event)
-		case getStatsChan := <-statsChan:
-			s.handleGetStats(getStatsChan)
-		case <-crashesChan:
-			s.handleCrash()
-		}
+func (m multiStats) Crash() {
+	for i := range m {
+		go m[i].Crash()
 	}
 }
 
-func (s *Stats) handleTraffic(evt trafficData) {
-	if evt.ingress {
-		s.Traffic.ingress += uint64(evt.traffic)
-	} else {
-		s.Traffic.egress += uint64(evt.traffic)
+func (m multiStats) AntiReplayDetected() {
+	for i := range m {
+		go m[i].AntiReplayDetected()
 	}
 }
 
-func (s *Stats) handleSpeed() {
-	s.Speed.ingress = s.Traffic.ingress - s.previousTraffic.ingress
-	s.Speed.egress = s.Traffic.egress - s.previousTraffic.egress
-	s.previousTraffic.ingress = s.Traffic.ingress
-	s.previousTraffic.egress = s.Traffic.egress
-}
+var S Stats
 
-func (s *Stats) handleConnection(evt connectionData) {
-	var inc uint32 = 1
-	if !evt.connected {
-		inc = ^uint32(0)
+func Init() error {
+	mux := http.NewServeMux()
+
+	instanceJSON := newStatsJSON(mux)
+	instancePrometheus, err := newStatsPrometheus(mux)
+	if err != nil {
+		return errors.Annotate(err, "Cannot initialize Prometheus")
 	}
 
-	var conn *connectionType
-	switch evt.connectionType {
-	case mtproto.ConnectionTypeAbridged:
-		conn = &s.Connections.Abridged
-	case mtproto.ConnectionTypeSecure:
-		conn = &s.Connections.Secure
-	default:
-		conn = &s.Connections.Intermediate
+	stats := []Stats{instanceJSON, instancePrometheus}
+	if config.C.StatsdStats.Addr.IP != nil {
+		instanceStatsd, err := newStatsStatsd()
+		if err != nil {
+			return errors.Annotate(err, "Cannot initialize StatsD")
+		}
+		stats = append(stats, instanceStatsd)
 	}
 
-	if evt.addr.IP.To4() != nil {
-		conn.IPv4 += inc
-	} else {
-		conn.IPv6 += inc
+	listener, err := net.Listen("tcp", config.C.StatsAddr.String())
+	if err != nil {
+		return errors.Annotate(err, "Cannot initialize stats server")
 	}
-}
 
-func (s *Stats) handleGetStats(getStatsChan chan<- Stats) {
-	getStatsChan <- *s
-}
+	srv := http.Server{
+		Handler: mux,
+	}
+	go srv.Serve(listener) // nolint: errcheck
 
-func (s *Stats) handleCrash() {
-	s.Crashes++
-}
+	S = multiStats(stats)
 
-// NewStats creates a new instance of Stats structure.
-func NewStats(conf *config.Config) *Stats {
-	return &Stats{
-		URLs:   conf.GetURLs(),
-		Uptime: uptime(time.Now()),
-	}
+	return nil
 }

newstats/stats_json.go → stats/stats_json.go

@@ -1,15 +1,16 @@
-package newstats
+package stats
 
 import (
 	"encoding/json"
 	"net"
 	"net/http"
+	"strconv"
 	"sync/atomic"
 	"time"
 
 	"go.uber.org/zap"
 
-	"github.com/9seconds/mtg/newprotocol"
+	"github.com/9seconds/mtg/conntypes"
 )
 
 type statsJSON struct {
@@ -51,7 +52,8 @@ type statsJSONTraffic struct {
 type statsJSONUptime time.Time
 
 func (s statsJSONUptime) MarshalJSON() ([]byte, error) {
-	return json.Marshal(time.Since(time.Time(s)).Seconds())
+	seconds := strconv.Itoa(int(time.Since(time.Time(s)).Seconds()))
+	return []byte(seconds), nil
 }
 
 func (s *statsJSON) IngressTraffic(traffic int) {
@@ -62,27 +64,27 @@ func (s *statsJSON) EgressTraffic(traffic int) {
 	atomic.AddUint64(&s.Traffic.Egress, uint64(traffic))
 }
 
-func (s *statsJSON) ClientConnected(connectionType newprotocol.ConnectionType, addr *net.TCPAddr) {
+func (s *statsJSON) ClientConnected(connectionType conntypes.ConnectionType, addr *net.TCPAddr) {
 	s.changeConnections(connectionType, addr, 1)
 }
 
-func (s *statsJSON) ClientDisconnected(connectionType newprotocol.ConnectionType, addr *net.TCPAddr) {
+func (s *statsJSON) ClientDisconnected(connectionType conntypes.ConnectionType, addr *net.TCPAddr) {
 	s.changeConnections(connectionType, addr, ^uint32(0))
 }
 
-func (s *statsJSON) changeConnections(connectionType newprotocol.ConnectionType, addr *net.TCPAddr, value uint32) {
+func (s *statsJSON) changeConnections(connectionType conntypes.ConnectionType, addr *net.TCPAddr, value uint32) {
 	var connections *statsJSONConnectionType
 
 	switch connectionType {
-	case newprotocol.ConnectionTypeAbridged:
+	case conntypes.ConnectionTypeAbridged:
 		connections = &s.Connections.Abridged
-	case newprotocol.ConnectionTypeSecure:
+	case conntypes.ConnectionTypeSecure:
 		connections = &s.Connections.Secured
 	default:
 		connections = &s.Connections.Intermediate
 	}
 
-	if addr.IP.To4() == nil {
+	if addr.IP.To4() != nil {
 		atomic.AddUint32(&connections.IPv4, value)
 	} else {
 		atomic.AddUint32(&connections.IPv6, value)
@@ -98,7 +100,9 @@ func (s *statsJSON) AntiReplayDetected() {
 }
 
 func newStatsJSON(mux *http.ServeMux) Stats {
-	instance := &statsJSON{}
+	instance := &statsJSON{
+		Uptime: statsJSONUptime(time.Now()),
+	}
 	logger := zap.S().Named("stats")
 
 	mux.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {

newstats/stats_prometheus.go → stats/stats_prometheus.go

@@ -1,4 +1,4 @@
-package newstats
+package stats
 
 import (
 	"net"
@@ -8,8 +8,8 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 
-	"github.com/9seconds/mtg/newconfig"
-	"github.com/9seconds/mtg/newprotocol"
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/conntypes"
 )
 
 type statsPrometheus struct {
@@ -27,23 +27,23 @@ func (s *statsPrometheus) EgressTraffic(traffic int) {
 	s.traffic.WithLabelValues("egress").Add(float64(traffic))
 }
 
-func (s *statsPrometheus) ClientConnected(connectionType newprotocol.ConnectionType, addr *net.TCPAddr) {
+func (s *statsPrometheus) ClientConnected(connectionType conntypes.ConnectionType, addr *net.TCPAddr) {
 	s.changeConnections(connectionType, addr, 1.0)
 }
 
-func (s *statsPrometheus) ClientDisconnected(connectionType newprotocol.ConnectionType, addr *net.TCPAddr) {
+func (s *statsPrometheus) ClientDisconnected(connectionType conntypes.ConnectionType, addr *net.TCPAddr) {
 	s.changeConnections(connectionType, addr, -1.0)
 }
 
-func (s *statsPrometheus) changeConnections(connectionType newprotocol.ConnectionType,
+func (s *statsPrometheus) changeConnections(connectionType conntypes.ConnectionType,
 	addr *net.TCPAddr,
 	increment float64) {
 	var labels [2]string
 
 	switch connectionType {
-	case newprotocol.ConnectionTypeAbridged:
+	case conntypes.ConnectionTypeAbridged:
 		labels[0] = "abridged"
-	case newprotocol.ConnectionTypeSecure:
+	case conntypes.ConnectionTypeSecure:
 		labels[0] = "secured"
 	default:
 		labels[0] = "intermediate"
@@ -69,22 +69,22 @@ func newStatsPrometheus(mux *http.ServeMux) (Stats, error) {
 	registry := prometheus.NewRegistry()
 	instance := &statsPrometheus{
 		connections: prometheus.NewGaugeVec(prometheus.GaugeOpts{
-			Namespace: newconfig.C.PrometheusStats.Prefix,
+			Namespace: config.C.PrometheusStats.Prefix,
 			Name:      "connections",
 			Help:      "Current number of connections to the proxy.",
 		}, []string{"type", "protocol"}),
 		traffic: prometheus.NewGaugeVec(prometheus.GaugeOpts{
-			Namespace: newconfig.C.PrometheusStats.Prefix,
+			Namespace: config.C.PrometheusStats.Prefix,
 			Name:      "traffic",
 			Help:      "Traffic passed through the proxy in bytes.",
 		}, []string{"direction"}),
 		crashes: prometheus.NewGauge(prometheus.GaugeOpts{
-			Namespace: newconfig.C.PrometheusStats.Prefix,
+			Namespace: config.C.PrometheusStats.Prefix,
 			Name:      "crashes",
 			Help:      "How many crashes happened.",
 		}),
 		antiReplays: prometheus.NewGauge(prometheus.GaugeOpts{
-			Namespace: newconfig.C.PrometheusStats.Prefix,
+			Namespace: config.C.PrometheusStats.Prefix,
 			Name:      "anti_replays",
 			Help:      "How many anti replay attacks were prevented.",
 		}),

newstats/stats_statsd.go → stats/stats_statsd.go

@@ -1,14 +1,14 @@
-package newstats
+package stats
 
 import (
 	"net"
 	"strings"
 
+	"github.com/juju/errors"
 	"gopkg.in/alexcesaro/statsd.v2"
 
-	"github.com/9seconds/mtg/newconfig"
-	"github.com/9seconds/mtg/newprotocol"
-	"github.com/juju/errors"
+	"github.com/9seconds/mtg/config"
+	"github.com/9seconds/mtg/conntypes"
 )
 
 type statsStatsd struct {
@@ -23,22 +23,22 @@ func (s *statsStatsd) EgressTraffic(traffic int) {
 	s.client.Count("traffic.egress", traffic)
 }
 
-func (s *statsStatsd) ClientConnected(connectionType newprotocol.ConnectionType, addr *net.TCPAddr) {
+func (s *statsStatsd) ClientConnected(connectionType conntypes.ConnectionType, addr *net.TCPAddr) {
 	s.changeConnections(connectionType, addr, 1)
 }
 
-func (s *statsStatsd) ClientDisconnected(connectionType newprotocol.ConnectionType, addr *net.TCPAddr) {
+func (s *statsStatsd) ClientDisconnected(connectionType conntypes.ConnectionType, addr *net.TCPAddr) {
 	s.changeConnections(connectionType, addr, -1)
 }
 
-func (s *statsStatsd) changeConnections(connectionType newprotocol.ConnectionType, addr *net.TCPAddr, value int) {
+func (s *statsStatsd) changeConnections(connectionType conntypes.ConnectionType, addr *net.TCPAddr, value int) {
 	var labels [3]string
 
 	labels[0] = "connections"
 	switch connectionType {
-	case newprotocol.ConnectionTypeAbridged:
+	case conntypes.ConnectionTypeAbridged:
 		labels[1] = "abridged"
-	case newprotocol.ConnectionTypeSecure:
+	case conntypes.ConnectionTypeSecure:
 		labels[1] = "secured"
 	default:
 		labels[1] = "intermediate"
@@ -62,15 +62,15 @@ func (s *statsStatsd) AntiReplayDetected() {
 
 func newStatsStatsd() (Stats, error) {
 	options := []statsd.Option{
-		statsd.Prefix(newconfig.C.StatsdStats.Prefix),
-		statsd.Network(newconfig.C.StatsdStats.Addr.Network()),
-		statsd.Address(newconfig.C.StatsdStats.Addr.String()),
-		statsd.TagsFormat(newconfig.C.StatsdStats.TagsFormat),
+		statsd.Prefix(config.C.StatsdStats.Prefix),
+		statsd.Network(config.C.StatsdStats.Addr.Network()),
+		statsd.Address(config.C.StatsdStats.Addr.String()),
+		statsd.TagsFormat(config.C.StatsdStats.TagsFormat),
 	}
 
-	if len(newconfig.C.StatsdStats.Tags) > 0 {
-		tags := make([]string, len(newconfig.C.StatsdStats.Tags)*2)
-		for k, v := range newconfig.C.StatsdStats.Tags {
+	if len(config.C.StatsdStats.Tags) > 0 {
+		tags := make([]string, len(config.C.StatsdStats.Tags)*2)
+		for k, v := range config.C.StatsdStats.Tags {
 			tags = append(tags, k, v)
 		}
 		options = append(options, statsd.Tags(tags...))

telegram/base.go

@@ -0,0 +1,70 @@
+package telegram
+
+import (
+	"context"
+	"math/rand"
+	"net"
+	"time"
+
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/utils"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+const telegramDialTimeout = 10 * time.Second
+
+type baseTelegram struct {
+	dialer net.Dialer
+
+	v4DefaultDC conntypes.DC
+	V6DefaultDC conntypes.DC
+	v4Addresses map[conntypes.DC][]string
+	v6Addresses map[conntypes.DC][]string
+}
+
+func (b *baseTelegram) dialToAddress(ctx context.Context,
+	cancel context.CancelFunc,
+	addr string) (wrappers.StreamReadWriteCloser, error) {
+	conn, err := b.dialer.Dial("tcp", addr)
+	if err != nil {
+		return nil, errors.Annotate(err, "Dial has failed")
+	}
+
+	if err := utils.InitTCP(conn); err != nil {
+		return nil, errors.Annotate(err, "Cannot initialize TCP socket")
+	}
+
+	return wrappers.NewTelegramConn(ctx, cancel, conn), nil
+}
+
+func (b *baseTelegram) dial(ctx context.Context,
+	cancel context.CancelFunc,
+	dc conntypes.DC,
+	protocol conntypes.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error) {
+	addr := ""
+
+	switch protocol {
+	case conntypes.ConnectionProtocolIPv4:
+		addr = b.chooseAddress(b.v4Addresses, dc, b.v4DefaultDC)
+	default:
+		addr = b.chooseAddress(b.v6Addresses, dc, b.V6DefaultDC)
+	}
+
+	return b.dialToAddress(ctx, cancel, addr)
+}
+
+func (b *baseTelegram) chooseAddress(addresses map[conntypes.DC][]string,
+	dc, defaultDC conntypes.DC) string {
+	addrs, ok := addresses[dc]
+	if !ok {
+		addrs, _ = addresses[defaultDC]
+	}
+
+	if len(addrs) > 0 {
+		return addrs[rand.Intn(len(addrs))]
+	}
+
+	return ""
+}

telegram/direct.go

@@ -4,28 +4,24 @@ import (
 	"context"
 	"net"
 
-	"github.com/juju/errors"
-
-	"github.com/9seconds/mtg/config"
-	"github.com/9seconds/mtg/mtproto"
-	"github.com/9seconds/mtg/obfuscated2"
+	"github.com/9seconds/mtg/conntypes"
 	"github.com/9seconds/mtg/wrappers"
 )
 
 const (
-	directV4DefaultIdx = 1
-	directV6DefaultIdx = 1
+	directV4DefaultIdx conntypes.DC = 1
+	directV6DefaultIdx conntypes.DC = 1
 )
 
 var (
-	directV4Addresses = map[int16][]string{
+	directV4Addresses = map[conntypes.DC][]string{
 		0: {"149.154.175.50:443"},
 		1: {"149.154.167.51:443"},
 		2: {"149.154.175.100:443"},
 		3: {"149.154.167.91:443"},
 		4: {"149.154.171.5:443"},
 	}
-	directV6Addresses = map[int16][]string{
+	directV6Addresses = map[conntypes.DC][]string{
 		0: {"[2001:b28:f23d:f001::a]:443"},
 		1: {"[2001:67c:04e8:f002::a]:443"},
 		2: {"[2001:b28:f23d:f003::a]:443"},
@@ -38,42 +34,28 @@ type directTelegram struct {
 	baseTelegram
 }
 
-func (t *directTelegram) Dial(ctx context.Context, cancel context.CancelFunc,
-	connID string, connOpts *mtproto.ConnectionOpts) (wrappers.StreamReadWriteCloser, error) {
-	dc := connOpts.DC
-	if dc < 0 {
+func (d *directTelegram) Dial(ctx context.Context,
+	cancel context.CancelFunc,
+	dc conntypes.DC,
+	protocol conntypes.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error) {
+	switch {
+	case dc < 0:
 		dc = -dc
-	} else if dc == 0 {
-		dc = 1
-	}
-
-	return t.baseTelegram.dial(ctx, cancel, dc-1, connID, connOpts.ConnectionProto)
-}
-
-func (t *directTelegram) Init(connOpts *mtproto.ConnectionOpts,
-	conn wrappers.StreamReadWriteCloser) (wrappers.Wrap, error) {
-	obfs2, frame := obfuscated2.MakeTelegramObfuscated2Frame(connOpts)
-
-	if _, err := conn.Write(frame); err != nil {
-		return nil, errors.Annotate(err, "Cannot write hadnshake frame")
+	case dc == 0:
+		dc = conntypes.DCDefaultIdx
 	}
 
-	return wrappers.NewStreamCipher(conn, obfs2.Encryptor, obfs2.Decryptor), nil
+	return d.baseTelegram.dial(ctx, cancel, dc-1, protocol)
 }
 
-// NewDirectTelegram returns Telegram instance which connects directly
-// to Telegram bypassing middleproxies.
-func NewDirectTelegram(conf *config.Config) Telegram {
+func NewDirectTelegram() Telegram {
 	return &directTelegram{
 		baseTelegram: baseTelegram{
-			dialer: tgDialer{
-				Dialer: net.Dialer{Timeout: telegramDialTimeout},
-				conf:   conf,
-			},
-			v4DefaultIdx: directV4DefaultIdx,
-			v6DefaultIdx: directV6DefaultIdx,
-			v4Addresses:  directV4Addresses,
-			v6Addresses:  directV6Addresses,
+			dialer:      net.Dialer{Timeout: telegramDialTimeout},
+			v4DefaultDC: directV4DefaultIdx,
+			V6DefaultDC: directV6DefaultIdx,
+			v4Addresses: directV4Addresses,
+			v6Addresses: directV6Addresses,
 		},
 	}
 }

telegram/interfaces.go

@@ -0,0 +1,15 @@
+package telegram
+
+import (
+	"context"
+
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/wrappers"
+)
+
+type Telegram interface {
+	Dial(context.Context,
+		context.CancelFunc,
+		conntypes.DC,
+		conntypes.ConnectionProtocol) (wrappers.StreamReadWriteCloser, error)
+}

utils/init_tcp.go

@@ -0,0 +1,25 @@
+package utils
+
+import (
+	"net"
+
+	"github.com/juju/errors"
+
+	"github.com/9seconds/mtg/config"
+)
+
+func InitTCP(conn net.Conn) error {
+	tcpConn := conn.(*net.TCPConn)
+
+	if err := tcpConn.SetNoDelay(true); err != nil {
+		return errors.Annotate(err, "Cannot set NO_DELAY")
+	}
+	if err := tcpConn.SetReadBuffer(config.C.BufferSize.Read); err != nil {
+		return errors.Annotate(err, "Cannot set read buffer size")
+	}
+	if err := tcpConn.SetWriteBuffer(config.C.BufferSize.Write); err != nil {
+		return errors.Annotate(err, "Cannot set write buffer size")
+	}
+
+	return nil
+}

utils/stream_cipher.go

@@ -0,0 +1,11 @@
+package utils
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+)
+
+func MakeStreamCipher(key, iv []byte) cipher.Stream {
+	block, _ := aes.NewCipher(key) // nolint: gosec
+	return cipher.NewCTR(block, iv)
+}

utils/uint24.go

@@ -1,15 +1,11 @@
 package utils
 
-// Uint24 is a replacement for the absent Go uint24 data type.
-// This data type is little endian.
 type Uint24 [3]byte
 
-// ToUint24 converts number to Uint24.
 func ToUint24(number uint32) Uint24 {
 	return Uint24{byte(number), byte(number >> 8), byte(number >> 16)}
 }
 
-// FromUint24 converts Uint24 to number.
 func FromUint24(number Uint24) uint32 {
 	return uint32(number[0]) + (uint32(number[1]) << 8) + (uint32(number[2]) << 16)
 }

newwrappers/interfaces.go → wrappers/interfaces.go

@@ -1,4 +1,4 @@
-package newwrappers
+package wrappers
 
 import (
 	"io"

newwrappers/wrapper_conn.go → wrappers/wrapper_conn.go

@@ -1,4 +1,4 @@
-package newwrappers
+package wrappers
 
 import (
 	"context"
@@ -10,7 +10,7 @@ import (
 	"github.com/juju/errors"
 	"go.uber.org/zap"
 
-	"github.com/9seconds/mtg/newconfig"
+	"github.com/9seconds/mtg/config"
 )
 
 const ConnIDLength = 8
@@ -126,11 +126,11 @@ func newConn(ctx context.Context,
 	localAddr := *parent.LocalAddr().(*net.TCPAddr)
 
 	if parent.RemoteAddr().(*net.TCPAddr).IP.To4() != nil {
-		if newconfig.C.PublicIPv4Addr.IP != nil {
-			localAddr.IP = newconfig.C.PublicIPv4Addr.IP
+		if config.C.PublicIPv4Addr.IP != nil {
+			localAddr.IP = config.C.PublicIPv4Addr.IP
 		}
-	} else if newconfig.C.PublicIPv6Addr.IP != nil {
-		localAddr.IP = newconfig.C.PublicIPv6Addr.IP
+	} else if config.C.PublicIPv6Addr.IP != nil {
+		localAddr.IP = config.C.PublicIPv6Addr.IP
 	}
 
 	logger := zap.S().With(
@@ -162,9 +162,8 @@ func NewClientConn(ctx context.Context,
 
 func NewTelegramConn(ctx context.Context,
 	cancel context.CancelFunc,
-	parent net.Conn,
-	connID ConnID) StreamReadWriteCloser {
-	return newConn(ctx, cancel, parent, connID, connPurposeTelegram)
+	parent net.Conn) StreamReadWriteCloser {
+	return newConn(ctx, cancel, parent, ConnID{}, connPurposeTelegram)
 }
 
 func NewConnID() ConnID {

newwrappers/wrapper_obfuscated2.go → wrappers/wrapper_obfuscated2.go

@@ -1,4 +1,4 @@
-package newwrappers
+package wrappers
 
 import (
 	"crypto/cipher"

newwrappers/wrapper_stats.go → wrappers/wrapper_stats.go

@@ -1,4 +1,4 @@
-package newwrappers
+package wrappers
 
 import (
 	"net"
@@ -6,7 +6,7 @@ import (
 
 	"go.uber.org/zap"
 
-	"github.com/9seconds/mtg/newstats"
+	"github.com/9seconds/mtg/stats"
 )
 
 type wrapperStats struct {
@@ -15,28 +15,28 @@ type wrapperStats struct {
 
 func (w *wrapperStats) Write(p []byte) (int, error) {
 	n, err := w.parent.Write(p)
-	newstats.S.EgressTraffic(n)
+	stats.S.EgressTraffic(n)
 
 	return n, err
 }
 
 func (w *wrapperStats) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
 	n, err := w.parent.WriteTimeout(p, timeout)
-	newstats.S.EgressTraffic(n)
+	stats.S.EgressTraffic(n)
 
 	return n, err
 }
 
 func (w *wrapperStats) Read(p []byte) (int, error) {
 	n, err := w.parent.Read(p)
-	newstats.S.IngressTraffic(n)
+	stats.S.IngressTraffic(n)
 
 	return n, err
 }
 
 func (w *wrapperStats) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
 	n, err := w.parent.ReadTimeout(p, timeout)
-	newstats.S.IngressTraffic(n)
+	stats.S.IngressTraffic(n)
 
 	return n, err
 }