Reworked base

go.mod

 	github.com/juju/loggo v0.0.0-20190526231331-6e530bcce5d8 // indirect
 	github.com/juju/testing v0.0.0-20190723135506-ce30eb24acd2 // indirect
 	github.com/kr/pretty v0.1.0 // indirect
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
+	github.com/modern-go/reflect2 v1.0.1 // indirect
 	github.com/pkg/errors v0.8.1 // indirect
 	github.com/prometheus/client_golang v1.1.0
 	github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4 // indirect

main.go

 package main
 
 import (
-	"encoding/json"
-	"fmt"
-	"io"
 	"math/rand"
 	"os"
 	"syscall"
 	"time"
 
 	"github.com/juju/errors"
-	"go.uber.org/zap"
-	"go.uber.org/zap/zapcore"
 	kingpin "gopkg.in/alecthomas/kingpin.v2"
 
-	"github.com/9seconds/mtg/config"
-	"github.com/9seconds/mtg/ntp"
-	"github.com/9seconds/mtg/proxy"
-	"github.com/9seconds/mtg/stats"
+	"github.com/9seconds/mtg/newcli"
+	"github.com/9seconds/mtg/newconfig"
 )
 
 var version = "dev" // this has to be set by build ld flags
 var (
 	app = kingpin.New("mtg", "Simple MTPROTO proxy.")
 
-	debug = app.Flag("debug",
+	generateSecretCommand = app.Command("generate-secret",
+		"Generate new secret")
+	generateSecretType = generateSecretCommand.Arg("type",
+		"A type of secret to generate. Valid options are 'simple', 'secured' and 'tls'").
+		Required().
+		Enum("simple", "secured", "tls")
+
+	proxyCommand = app.Command("proxy",
+		"Run new proxy instance")
+	proxyDebug = proxyCommand.Flag("debug",
 		"Run in debug mode.").
 		Short('d').
 		Envar("MTG_DEBUG").
 		Bool()
-	verbose = app.Flag("verbose",
+	proxyVerbose = proxyCommand.Flag("verbose",
 		"Run in verbose mode.").
 		Short('v').
 		Envar("MTG_VERBOSE").
 		Bool()
-
-	bindIP = app.Flag("bind-ip",
+	proxyBindIP = proxyCommand.Flag("bind-ip",
 		"Which IP to bind to.").
 		Short('b').
 		Envar("MTG_IP").
 		Default("127.0.0.1").
 		IP()
-	bindPort = app.Flag("bind-port",
+	proxyBindPort = proxyCommand.Flag("bind-port",
 		"Which port to bind to.").
 		Short('p').
 		Envar("MTG_PORT").
 		Default("3128").
 		Uint16()
-
-	publicIPv4 = app.Flag("public-ipv4",
+	proxyPublicIPv4 = proxyCommand.Flag("public-ipv4",
 		"Which IPv4 address is public.").
 		Short('4').
 		Envar("MTG_IPV4").
 		IP()
-	publicIPv4Port = app.Flag("public-ipv4-port",
+	proxyPublicIPv4Port = proxyCommand.Flag("public-ipv4-port",
 		"Which IPv4 port is public. Default is 'bind-port' value.").
 		Envar("MTG_IPV4_PORT").
 		Uint16()
-
-	publicIPv6 = app.Flag("public-ipv6",
+	proxyPublicIPv6 = proxyCommand.Flag("public-ipv6",
 		"Which IPv6 address is public.").
 		Short('6').
 		Envar("MTG_IPV6").
 		IP()
-	publicIPv6Port = app.Flag("public-ipv6-port",
+	proxyPublicIPv6Port = proxyCommand.Flag("public-ipv6-port",
 		"Which IPv6 port is public. Default is 'bind-port' value.").
 		Envar("MTG_IPV6_PORT").
 		Uint16()
-
-	statsIP = app.Flag("stats-ip",
+	proxyStatsIP = proxyCommand.Flag("stats-ip",
 		"Which IP bind stats server to.").
 		Short('t').
 		Envar("MTG_STATS_IP").
 		Default("127.0.0.1").
 		IP()
-	statsPort = app.Flag("stats-port",
+	proxyStatsPort = proxyCommand.Flag("stats-port",
 		"Which port bind stats to.").
 		Short('q').
 		Envar("MTG_STATS_PORT").
 		Default("3129").
 		Uint16()
-
-	statsdIP = app.Flag("statsd-ip",
+	proxyStatsdIP = proxyCommand.Flag("statsd-ip",
 		"Which IP should we use for working with statsd.").
 		Envar("MTG_STATSD_IP").
-		String()
-	statsdPort = app.Flag("statsd-port",
+		IP()
+	proxyStatsdPort = proxyCommand.Flag("statsd-port",
 		"Which port should we use for working with statsd.").
 		Envar("MTG_STATSD_PORT").
 		Default("8125").
 		Uint16()
-	statsdNetwork = app.Flag("statsd-network",
+	proxyStatsdNetwork = proxyCommand.Flag("statsd-network",
 		"Which network is used to work with statsd. Only 'tcp' and 'udp' are supported.").
 		Envar("MTG_STATSD_NETWORK").
 		Default("udp").
-		String()
-	statsdPrefix = app.Flag("statsd-prefix",
+		Enum("udp", "tcp")
+	proxyStatsdPrefix = proxyCommand.Flag("statsd-prefix",
 		"Which bucket prefix should we use for sending stats to statsd.").
 		Envar("MTG_STATSD_PREFIX").
 		Default("mtg").
 		String()
-	statsdTagsFormat = app.Flag("statsd-tags-format",
+	proxyStatsdTagsFormat = proxyCommand.Flag("statsd-tags-format",
 		"Which tag format should we use to send stats metrics. Valid options are 'datadog' and 'influxdb'.").
 		Envar("MTG_STATSD_TAGS_FORMAT").
-		String()
-	statsdTags = app.Flag("statsd-tags",
+		Default("influxdb").
+		Enum("datadog", "influxdb")
+	proxyStatsdTags = proxyCommand.Flag("statsd-tags",
 		"Tags to use for working with statsd (specified as 'key=value').").
 		Envar("MTG_STATSD_TAGS").
 		StringMap()
-
-	prometheusPrefix = app.Flag("prometheus-prefix",
+	proxyPrometheusPrefix = proxyCommand.Flag("prometheus-prefix",
 		"Which namespace to use to send stats to Prometheus.").
 		Envar("MTG_PROMETHEUS_PREFIX").
 		Default("mtg").
 		String()
-
-	writeBufferSize = app.Flag("write-buffer",
+	proxyWriteBufferSize = proxyCommand.Flag("write-buffer",
 		"Write buffer size in bytes. You can think about it as a buffer from client to Telegram.").
 		Short('w').
 		Envar("MTG_BUFFER_WRITE").
 		Default("65536").
 		Uint32()
-	readBufferSize = app.Flag("read-buffer",
+	proxyReadBufferSize = proxyCommand.Flag("read-buffer",
 		"Read buffer size in bytes. You can think about it as a buffer from Telegram to client.").
 		Short('r').
 		Envar("MTG_BUFFER_READ").
 		Default("131072").
 		Uint32()
-	secureOnly = app.Flag("secure-only",
-		"Support clients with dd-secrets only.").
-		Short('s').
-		Envar("MTG_SECURE_ONLY").
-		Bool()
-
-	antiReplayMaxSize = app.Flag("anti-replay-max-size",
+	proxyAntiReplayMaxSize = proxyCommand.Flag("anti-replay-max-size",
 		"Max size of antireplay cache in megabytes.").
 		Envar("MTG_ANTIREPLAY_MAXSIZE").
 		Default("128").
 		Int()
-	antiReplayEvictionTime = app.Flag("anti-replay-eviction-time",
+	proxyAntiReplayEvictionTime = proxyCommand.Flag("anti-replay-eviction-time",
 		"Eviction time period for obfuscated2 handshakes").
 		Envar("MTG_ANTIREPLAY_EVICTIONTIME").
 		Default("168h").
 		Duration()
-
-	secret = app.Arg("secret", "Secret of this proxy.").Required().HexBytes()
-	adtag  = app.Arg("adtag", "ADTag of the proxy.").HexBytes()
+	proxySecret = proxyCommand.Arg("secret", "Secret of this proxy.").Required().HexBytes()
+	proxyAdtag  = proxyCommand.Arg("adtag", "ADTag of the proxy.").HexBytes()
 )
 
-func main() { // nolint: gocyclo
+func main() {
 	rand.Seed(time.Now().UTC().UnixNano())
 	app.Version(version)
 	app.HelpFlag.Short('h')
 
-	kingpin.MustParse(app.Parse(os.Args[1:]))
-
-	err := setRLimit()
-	if err != nil {
-		usage(err.Error())
-	}
-
-	conf, err := config.NewConfig(*debug, *verbose,
-		*writeBufferSize, *readBufferSize,
-		*bindIP, *publicIPv4, *publicIPv6, *statsIP,
-		*bindPort, *publicIPv4Port, *publicIPv6Port, *statsPort, *statsdPort,
-		*statsdIP, *statsdNetwork, *statsdPrefix, *statsdTagsFormat,
-		*statsdTags, *prometheusPrefix, *secureOnly,
-		*antiReplayMaxSize, *antiReplayEvictionTime,
-		*secret, *adtag,
-	)
-	if err != nil {
-		usage(err.Error())
+	if err := setRLimit(); err != nil {
+		newcli.Fatal(err.Error())
 	}
 
-	atom := zap.NewAtomicLevel()
-	switch {
-	case conf.Debug:
-		atom.SetLevel(zapcore.DebugLevel)
-	case conf.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
-
-	printURLs(conf.GetURLs())
-	zap.S().Debugw("Configuration", "config", conf)
-
-	if conf.UseMiddleProxy() {
-		zap.S().Infow("Use middle proxy connection to Telegram")
-		if diff, err := ntp.Fetch(); err != nil {
-			zap.S().Warnw("Could not fetch time data from NTP")
-		} else {
-			if diff >= time.Second {
-				usage(fmt.Sprintf("You choose to use middle proxy but your clock drift (%s) "+
-					"is bigger than 1 second. Please, sync your time", diff))
-			}
-			go ntp.AutoUpdate()
+	switch kingpin.MustParse(app.Parse(os.Args[1:])) {
+	case generateSecretCommand.FullCommand():
+		newcli.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},
+		)
+		if err != nil {
+			newcli.Fatal(err.Error())
 		}
-	} else {
-		zap.S().Infow("Use direct connection to Telegram")
-	}
-
-	if err := stats.Init(conf); err != nil {
-		panic(err)
-	}
 
-	server, err := proxy.NewProxy(conf)
-	if err != nil {
-		panic(err)
-	}
-	if err := server.Serve(); err != nil {
-		zap.S().Fatalw("Server stopped", "error", err)
+		if err := newcli.Proxy(); err != nil {
+			newcli.Fatal(err.Error())
+		}
 	}
 }
 
 
 	return
 }
-
-func printURLs(data interface{}) {
-	encoder := json.NewEncoder(os.Stdout)
-	encoder.SetEscapeHTML(false)
-	encoder.SetIndent("", "  ")
-
-	err := encoder.Encode(data)
-	if err != nil {
-		panic(err)
-	}
-}
-
-func usage(msg string) {
-	io.WriteString(os.Stderr, msg+"\n") // nolint: errcheck, gosec
-	os.Exit(1)
-}

newantireplay/cache.go

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

newantireplay/hasher.go

+package newantireplay
+
+import "github.com/cespare/xxhash"
+
+type hasher struct{}
+
+func (h hasher) Sum64(value string) uint64 {
+	return xxhash.Sum64String(value)
+}

newcli/generate.go

+package newcli
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+
+	"github.com/9seconds/mtg/newconfig"
+)
+
+func Generate(secretType string) {
+	data := make([]byte, newconfig.SimpleSecretLength)
+	if _, err := rand.Read(data); err != nil {
+		panic(err)
+	}
+	secret := hex.EncodeToString(data)
+
+	switch secretType {
+	case "simple":
+		PrintStdout(secret)
+	case "secured":
+		PrintStdout("dd" + secret)
+	default:
+		Fatal("Unknown secret type " + secret)
+	}
+}

newcli/proxy.go

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

newcli/utils.go

+package newcli
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+)
+
+func Fatal(args ...interface{}) {
+	PrintStderr(args...)
+	os.Exit(1)
+}
+
+func PrintStderr(args ...interface{}) {
+	fmt.Fprintln(os.Stderr, args...)
+}
+
+func PrintStdout(args ...interface{}) {
+	fmt.Println(args...)
+}
+
+func PrintJSONStderr(data interface{}) {
+	printJSON(os.Stderr, data)
+}
+
+func PrintJSONStdout(data interface{}) {
+	printJSON(os.Stdout, data)
+}
+
+func printJSON(writer io.Writer, data interface{}) {
+	encoder := json.NewEncoder(writer)
+	encoder.SetEscapeHTML(false)
+	encoder.SetIndent("", "  ")
+
+	if err := encoder.Encode(data); err != nil {
+		panic(err)
+	}
+}

config2/config.go → newconfig/config.go

-package config2
+package newconfig
 
 import (
 	"bytes"
-	"context"
 	"encoding/json"
 	"net"
 	"strconv"
-	"sync"
 	"time"
 
 	"github.com/juju/errors"
+	"go.uber.org/zap"
 	statsd "gopkg.in/alexcesaro/statsd.v2"
 )
 
-type SecretType byte
+type SecretMode uint8
 
-func (s SecretType) String() string {
+func (s SecretMode) String() string {
 	switch s {
-	case SecretTypeMain:
-		return "main"
-	case SecretTypeSecured:
+	case SecretModeSimple:
+		return "simple"
+	case SecretModeSecured:
 		return "secured"
-	default:
-		return "tls"
 	}
+	return "tls"
 }
 
 const (
-	SecretTypeMain = 1 << iota
-	SecretTypeSecured
-	SecretTypeTLS
+	SecretModeSimple SecretMode = iota
+	SecretModeSecured
+	SecretModeTLS
 )
 
-const (
-	FlagDebug   = "debug"
-	FlagVerbose = "verbose"
-
-	FlagBindIP         = "bind-ip"
-	FlagBindPort       = "bind-port"
-	FlagPublicIPv4     = "public-ipv4"
-	FlagPublicIPv4Port = "public-ipv4-port"
-	FlagPublicIPv6     = "public-ipv6"
-	FlagPublicIPv6Port = "public-ipv6-port"
-	FlagStatsIP        = "stats-ip"
-	FlagStatsPort      = "stats-port"
-
-	FlagStatsdIP         = "statsd-ip"
-	FlagStatsdPort       = "statsd-port"
-	FlagStatsdNetwork    = "statsd-network"
-	FlagStatsdPrefix     = "statsd-prefix"
-	FlagStatsdTagsFormat = "statsd-tags-format"
-	FlagStatsdTags       = "statsd-tags"
-
-	FlagPrometheusPrefix = "prometheus-prefix"
+const SimpleSecretLength = 16
 
-	FlagWriteBufferSize = "write-buffer"
-	FlagReadBufferSize  = "read-buffer"
+type OptionType uint8
 
-	FlagSecureOnly = "secure-only"
-
-	FlagAntiReplayMaxSize      = "anti-replay-max-size"
-	FlagAntiReplayEvictionTime = "anti-replay-eviction-time"
-
-	FlagSecret = "secret"
-	FlagAdtag  = "adtag"
+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 {
 }
 
 type Stats struct {
-	Prefix  string `json:"prefix"`
-	Enabled bool   `json:"enabled"`
+	Prefix string `json:"prefix"`
 }
 
 type StatsdStats struct {
 
 	Debug      bool       `json:"debug"`
 	Verbose    bool       `json:"verbose"`
-	SecureOnly bool       `json:"secure_only"`
-	SecretType SecretType `json:"secret_type"`
+	SecretMode SecretMode `json:"secret_mode"`
 	Secret     []byte     `json:"secret"`
 	AdTag      []byte     `json:"adtag"`
 }
 }
 
 type ConfigOpt struct {
-	Name  string
-	Value interface{}
+	Option OptionType
+	Value  interface{}
 }
 
 var C = Config{}
 
 func Init(options ...ConfigOpt) error { // nolint: gocyclo
 	for _, opt := range options {
-		switch opt.Name {
-		case FlagDebug:
+		switch opt.Option {
+		case OptionTypeDebug:
 			C.Debug = opt.Value.(bool)
-		case FlagVerbose:
+		case OptionTypeVerbose:
 			C.Verbose = opt.Value.(bool)
-		case FlagBindIP:
+		case OptionTypeBindIP:
 			C.ListenAddr.IP = opt.Value.(net.IP)
-		case FlagBindPort:
-			C.ListenAddr.Port = opt.Value.(int)
-		case FlagPublicIPv4:
+		case OptionTypeBindPort:
+			C.ListenAddr.Port = int(opt.Value.(uint16))
+		case OptionTypePublicIPv4:
 			C.PublicIPv4Addr.IP = opt.Value.(net.IP)
-		case FlagPublicIPv4Port:
-			C.PublicIPv4Addr.Port = opt.Value.(int)
-		case FlagPublicIPv6:
+		case OptionTypePublicIPv4Port:
+			C.PublicIPv4Addr.Port = int(opt.Value.(uint16))
+		case OptionTypePublicIPv6:
 			C.PublicIPv6Addr.IP = opt.Value.(net.IP)
-		case FlagPublicIPv6Port:
-			C.PublicIPv6Addr.Port = opt.Value.(int)
-		case FlagStatsIP:
+		case OptionTypePublicIPv6Port:
+			C.PublicIPv6Addr.Port = int(opt.Value.(uint16))
+		case OptionTypeStatsIP:
 			C.StatsAddr.IP = opt.Value.(net.IP)
-		case FlagStatsPort:
-			C.StatsAddr.Port = opt.Value.(int)
-		case FlagStatsdIP:
+		case OptionTypeStatsPort:
+			C.StatsAddr.Port = int(opt.Value.(uint16))
+		case OptionTypeStatsdIP:
 			C.StatsdStats.Addr.IP = opt.Value.(net.IP)
-		case FlagStatsdPort:
-			C.StatsdStats.Addr.Port = opt.Value.(int)
-		case FlagStatsdNetwork:
+		case OptionTypeStatsdPort:
+			C.StatsdStats.Addr.Port = int(opt.Value.(uint16))
+		case OptionTypeStatsdNetwork:
 			C.StatsdStats.Addr.net = opt.Value.(string)
-		case FlagStatsdPrefix:
+		case OptionTypeStatsdPrefix:
 			C.StatsdStats.Prefix = opt.Value.(string)
-		case FlagStatsdTagsFormat:
+		case OptionTypeStatsdTagsFormat:
 			value := opt.Value.(string)
 			switch value {
 			case "datadog":
 			default:
 				return errors.Errorf("Incorrect statsd tag %s", value)
 			}
-		case FlagStatsdTags:
+		case OptionTypeStatsdTags:
 			C.StatsdStats.Tags = opt.Value.(map[string]string)
-		case FlagPrometheusPrefix:
+		case OptionTypePrometheusPrefix:
 			C.PrometheusStats.Prefix = opt.Value.(string)
-		case FlagWriteBufferSize:
-			C.BufferSize.Write = opt.Value.(int)
-		case FlagReadBufferSize:
-			C.BufferSize.Read = opt.Value.(int)
-		case FlagAntiReplayMaxSize:
+		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 FlagAntiReplayEvictionTime:
+		case OptionTypeAntiReplayEvictionTime:
 			C.AntiReplay.EvictionTime = opt.Value.(time.Duration)
-		case FlagSecureOnly:
-			C.SecureOnly = opt.Value.(bool)
-		case FlagSecret:
+		case OptionTypeSecret:
 			C.Secret = opt.Value.([]byte)
-		case FlagAdtag:
+		case OptionTypeAdtag:
 			C.AdTag = opt.Value.([]byte)
+		default:
+			return errors.Errorf("Unknown tag %v", opt.Option)
 		}
 	}
 
-	var defaultStatsdTags statsd.TagFormat
-	if C.StatsdStats.TagsFormat == defaultStatsdTags {
-		C.StatsdStats.TagsFormat = statsd.Datadog
-	}
-	if C.StatsdStats.Addr.net == "" {
-		C.StatsdStats.Addr.net = "udp"
-	}
-
 	switch {
-	case len(C.Secret) == 17 && bytes.HasPrefix(C.Secret, []byte{0xdd}):
-		C.SecretType = SecretTypeSecured
+	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) == 16:
-		C.SecretType = SecretTypeMain
+	case len(C.Secret) == SimpleSecretLength:
+		C.SecretMode = SecretModeSimple
 	default:
 		return errors.New("Incorrect secret")
 	}
 		C.PublicIPv6Addr.Port = C.ListenAddr.Port
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-	wg := &sync.WaitGroup{}
-	done := make(chan struct{})
-
+	foundAddress := C.PublicIPv4Addr.IP != nil || C.PublicIPv6Addr.IP != nil
 	if C.PublicIPv4Addr.IP == nil {
-		wg.Add(1)
-		go func() {
-			getGlobalIPv4(ctx, cancel)
-			wg.Done()
-		}()
+		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 {
-		wg.Add(1)
-		go func() {
-			getGlobalIPv6(ctx, cancel)
-			wg.Done()
-
-		}()
+		ip, err := getGlobalIPv6()
+		if err != nil {
+			zap.S().Warnw("Cannot resolve public address", "error", err)
+		} else {
+			C.PublicIPv6Addr.IP = ip
+			foundAddress = true
+		}
 	}
-	go func() {
-		wg.Wait()
-		close(done)
-	}()
-
-	select {
-	case <-done:
-		return nil
-	case <-ctx.Done():
-		return ctx.Err()
+
+	if !foundAddress {
+		return errors.New("Cannot resolve any public address")
 	}
+
+	return nil
 }

config2/global_ips.go → newconfig/global_ips.go

-package config2
+package newconfig
 
 import (
 	"context"
 	"time"
 
 	"github.com/juju/errors"
-	"go.uber.org/zap"
 )
 
 const (
 	ifconfigTimeout = 10 * time.Second
 )
 
-func getGlobalIPv4(ctx context.Context, cancel context.CancelFunc) {
-	ip, err := fetchIP(ctx, "tcp4")
+func getGlobalIPv4() (net.IP, error) {
+	ip, err := fetchIP("tcp4")
 	if err != nil || ip.To4() == nil {
-		cancel()
-		zap.S().Errorw("Cannot find public ipv4 address", "error", err)
-		return
+		return nil, errors.Annotate(err, "Cannot find public ipv4 address")
 	}
-	C.PublicIPv4Addr.IP = ip
+	return ip, nil
 }
 
-func getGlobalIPv6(ctx context.Context, cancel context.CancelFunc) {
-	ip, err := fetchIP(ctx, "tcp6")
+func getGlobalIPv6() (net.IP, error) {
+	ip, err := fetchIP("tcp6")
 	if err != nil || ip.To4() != nil {
-		cancel()
-		zap.S().Errorw("Cannot find public ipv6 address", "error", err)
-		return
+		return nil, errors.Annotate(err, "Cannot find public ipv6 address")
 	}
-	C.PublicIPv6Addr.IP = ip
+	return ip, nil
 }
 
-func fetchIP(ctx context.Context, network string) (net.IP, error) {
+func fetchIP(network string) (net.IP, error) {
 	dialer := &net.Dialer{FallbackDelay: -1}
 	client := &http.Client{
 		Jar:     nil,
 		},
 	}
 
-	req, err := http.NewRequest("GET", ifconfigAddress, nil)
-	if err != nil {
-		panic(err)
-	}
-
-	resp, err := client.Do(req.WithContext(ctx))
+	resp, err := client.Get(ifconfigAddress)
 	if err != nil {
-		if resp.Body != nil {
+		if resp != nil {
 			io.Copy(ioutil.Discard, resp.Body) // nolint: errcheck
 		}
 		return nil, errors.Annotate(err, "Cannot perform a request")

config2/urls.go → newconfig/urls.go

-package config2
+package newconfig
 
 import (
 	"encoding/hex"
 
 func GetURLs() (urls IPURLs) {
 	secret := ""
-	switch C.SecretType {
-	case SecretTypeMain, SecretTypeSecured:
+	switch C.SecretMode {
+	case SecretModeSimple:
 		secret = hex.EncodeToString(C.Secret)
-		if C.SecureOnly {
-			secret = "dd" + secret
-		}
+	case SecretModeSecured:
+		secret = "dd" + hex.EncodeToString(C.Secret)
 	}
 
 	urls.IPv4 = makeURLs(&C.PublicIPv4Addr, secret)

newobfuscated2/client_protocol.go

+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

+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

+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

+package newprotocol
+
+type BaseProtocol struct {
+	ConnectionType     ConnectionType
+	ConnectionProtocol ConnectionProtocol
+	DC                 int16
+}

newprotocol/connection_protocol.go

+package newprotocol
+
+type ConnectionProtocol uint8
+
+func (c ConnectionProtocol) String() string {
+	switch c {
+	case ConnectionProtocolAny:
+		return "any"
+	case ConnectionProtocolIPv4:
+		return "ipv4"
+	}
+	return "ipv6"
+}
+
+const (
+	ConnectionProtocolIPv4 ConnectionProtocol = 1
+	ConnectionProtocolIPv6                    = ConnectionProtocolIPv4 << 1
+	ConnectionProtocolAny                     = ConnectionProtocolIPv4 | ConnectionProtocolIPv6
+)

newprotocol/connection_type.go

+package newprotocol
+
+type ConnectionType uint8
+
+const (
+	ConnectionTypeUnknown ConnectionType = iota
+	ConnectionTypeAbridged
+	ConnectionTypeIntermediate
+	ConnectionTypeSecure
+)
+
+var (
+	ConnectionTagAbridged     = []byte{0xef, 0xef, 0xef, 0xef}
+	ConnectionTagIntermediate = []byte{0xee, 0xee, 0xee, 0xee}
+	ConnectionTagSecure       = []byte{0xdd, 0xdd, 0xdd, 0xdd}
+)
+
+func (t ConnectionType) Tag() []byte {
+	switch t {
+	case ConnectionTypeAbridged:
+		return ConnectionTagAbridged
+	case ConnectionTypeIntermediate:
+		return ConnectionTagIntermediate
+	default:
+		return ConnectionTagSecure
+	}
+}

newproxy/proxy.go

+package newproxy

newstats/stats.go

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

newstats/stats_json.go

+package newstats
+
+import (
+	"encoding/json"
+	"net"
+	"net/http"
+	"sync/atomic"
+	"time"
+
+	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/newprotocol"
+)
+
+type statsJSON struct {
+	Connections statsJSONConnections `json:"connections"`
+	Traffic     statsJSONTraffic     `json:"traffic"`
+	Uptime      statsJSONUptime      `json:"uptime"`
+	Crashes     uint32               `json:"crashes"`
+	AntiReplays uint32               `json:"anti_replay_detected"`
+}
+
+type statsBaseJSONConnections struct {
+	All          statsJSONConnectionType `json:"all"`
+	Abridged     statsJSONConnectionType `json:"abridged"`
+	Intermediate statsJSONConnectionType `json:"intermediate"`
+	Secured      statsJSONConnectionType `json:"secured"`
+}
+
+type statsJSONConnections struct {
+	statsBaseJSONConnections
+}
+
+type statsJSONConnectionType struct {
+	IPv4 uint32 `json:"ipv4"`
+	IPv6 uint32 `json:"ipv6"`
+}
+
+func (c statsJSONConnections) MarshalJSON() ([]byte, error) {
+	c.All.IPv4 = c.Abridged.IPv4 + c.Intermediate.IPv4 + c.Secured.IPv4
+	c.All.IPv6 = c.Abridged.IPv6 + c.Intermediate.IPv6 + c.Secured.IPv6
+
+	return json.Marshal(c.statsBaseJSONConnections)
+}
+
+type statsJSONTraffic struct {
+	Ingress uint64 `json:"ingress"`
+	Egress  uint64 `json:"egress"`
+}
+
+type statsJSONUptime time.Time
+
+func (s statsJSONUptime) MarshalJSON() ([]byte, error) {
+	return json.Marshal(time.Since(time.Time(s)).Seconds())
+}
+
+func (s *statsJSON) IngressTraffic(traffic int) {
+	atomic.AddUint64(&s.Traffic.Ingress, uint64(traffic))
+}
+
+func (s *statsJSON) EgressTraffic(traffic int) {
+	atomic.AddUint64(&s.Traffic.Egress, uint64(traffic))
+}
+
+func (s *statsJSON) ClientConnected(connectionType newprotocol.ConnectionType, addr *net.TCPAddr) {
+	s.changeConnections(connectionType, addr, 1)
+}
+
+func (s *statsJSON) ClientDisconnected(connectionType newprotocol.ConnectionType, addr *net.TCPAddr) {
+	s.changeConnections(connectionType, addr, ^uint32(0))
+}
+
+func (s *statsJSON) changeConnections(connectionType newprotocol.ConnectionType, addr *net.TCPAddr, value uint32) {
+	var connections *statsJSONConnectionType
+
+	switch connectionType {
+	case newprotocol.ConnectionTypeAbridged:
+		connections = &s.Connections.Abridged
+	case newprotocol.ConnectionTypeSecure:
+		connections = &s.Connections.Secured
+	default:
+		connections = &s.Connections.Intermediate
+	}
+
+	if addr.IP.To4() == nil {
+		atomic.AddUint32(&connections.IPv4, value)
+	} else {
+		atomic.AddUint32(&connections.IPv6, value)
+	}
+}
+
+func (s *statsJSON) Crash() {
+	atomic.AddUint32(&s.Crashes, 1)
+}
+
+func (s *statsJSON) AntiReplayDetected() {
+	atomic.AddUint32(&s.AntiReplays, 1)
+}
+
+func newStatsJSON(mux *http.ServeMux) Stats {
+	instance := &statsJSON{}
+	logger := zap.S().Named("stats")
+
+	mux.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		first, err := json.Marshal(instance)
+		if err != nil {
+			logger.Errorw("Cannot encode json", "error", err)
+			http.Error(w, "Internal server error", http.StatusServiceUnavailable)
+			return
+		}
+
+		interim := map[string]interface{}{}
+		if err := json.Unmarshal(first, &interim); err != nil {
+			panic(err)
+		}
+
+		encoder := json.NewEncoder(w)
+		encoder.SetEscapeHTML(false)
+		encoder.SetIndent("", "  ")
+		if err := encoder.Encode(interim); err != nil {
+			logger.Errorw("Cannot encode json", "error", err)
+		}
+	})
+
+	return instance
+}

newstats/stats_prometheus.go

+package newstats
+
+import (
+	"net"
+	"net/http"
+
+	"github.com/juju/errors"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+
+	"github.com/9seconds/mtg/newconfig"
+	"github.com/9seconds/mtg/newprotocol"
+)
+
+type statsPrometheus struct {
+	connections *prometheus.GaugeVec
+	traffic     *prometheus.GaugeVec
+	crashes     prometheus.Gauge
+	antiReplays prometheus.Gauge
+}
+
+func (s *statsPrometheus) IngressTraffic(traffic int) {
+	s.traffic.WithLabelValues("ingress").Add(float64(traffic))
+}
+
+func (s *statsPrometheus) EgressTraffic(traffic int) {
+	s.traffic.WithLabelValues("egress").Add(float64(traffic))
+}
+
+func (s *statsPrometheus) ClientConnected(connectionType newprotocol.ConnectionType, addr *net.TCPAddr) {
+	s.changeConnections(connectionType, addr, 1.0)
+}
+
+func (s *statsPrometheus) ClientDisconnected(connectionType newprotocol.ConnectionType, addr *net.TCPAddr) {
+	s.changeConnections(connectionType, addr, -1.0)
+}
+
+func (s *statsPrometheus) changeConnections(connectionType newprotocol.ConnectionType,
+	addr *net.TCPAddr,
+	increment float64) {
+	var labels [2]string
+
+	switch connectionType {
+	case newprotocol.ConnectionTypeAbridged:
+		labels[0] = "abridged"
+	case newprotocol.ConnectionTypeSecure:
+		labels[0] = "secured"
+	default:
+		labels[0] = "intermediate"
+	}
+
+	labels[1] = "ipv4"
+	if addr.IP.To4() == nil {
+		labels[1] = "ipv6"
+	}
+
+	s.connections.WithLabelValues(labels[:]...).Add(increment)
+}
+
+func (s *statsPrometheus) Crash() {
+	s.crashes.Inc()
+}
+
+func (s *statsPrometheus) AntiReplayDetected() {
+	s.antiReplays.Inc()
+}
+
+func newStatsPrometheus(mux *http.ServeMux) (Stats, error) {
+	registry := prometheus.NewRegistry()
+	instance := &statsPrometheus{
+		connections: prometheus.NewGaugeVec(prometheus.GaugeOpts{
+			Namespace: newconfig.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,
+			Name:      "traffic",
+			Help:      "Traffic passed through the proxy in bytes.",
+		}, []string{"direction"}),
+		crashes: prometheus.NewGauge(prometheus.GaugeOpts{
+			Namespace: newconfig.C.PrometheusStats.Prefix,
+			Name:      "crashes",
+			Help:      "How many crashes happened.",
+		}),
+		antiReplays: prometheus.NewGauge(prometheus.GaugeOpts{
+			Namespace: newconfig.C.PrometheusStats.Prefix,
+			Name:      "anti_replays",
+			Help:      "How many anti replay attacks were prevented.",
+		}),
+	}
+
+	if err := registry.Register(instance.connections); err != nil {
+		return nil, errors.Annotate(err, "Cannot register metrics for connections")
+	}
+	if err := registry.Register(instance.traffic); err != nil {
+		return nil, errors.Annotate(err, "Cannot register metrics for traffic")
+	}
+	if err := registry.Register(instance.crashes); err != nil {
+		return nil, errors.Annotate(err, "Cannot register metrics for crashes")
+	}
+	if err := registry.Register(instance.antiReplays); err != nil {
+		return nil, errors.Annotate(err, "Cannot register metrics for anti replays")
+	}
+
+	handler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{})
+	mux.Handle("/prometheus", handler)
+
+	return instance, nil
+}

newstats/stats_statsd.go

+package newstats
+
+import (
+	"net"
+	"strings"
+
+	"gopkg.in/alexcesaro/statsd.v2"
+
+	"github.com/9seconds/mtg/newconfig"
+	"github.com/9seconds/mtg/newprotocol"
+	"github.com/juju/errors"
+)
+
+type statsStatsd struct {
+	client *statsd.Client
+}
+
+func (s *statsStatsd) IngressTraffic(traffic int) {
+	s.client.Count("traffic.ingress", traffic)
+}
+
+func (s *statsStatsd) EgressTraffic(traffic int) {
+	s.client.Count("traffic.egress", traffic)
+}
+
+func (s *statsStatsd) ClientConnected(connectionType newprotocol.ConnectionType, addr *net.TCPAddr) {
+	s.changeConnections(connectionType, addr, 1)
+}
+
+func (s *statsStatsd) ClientDisconnected(connectionType newprotocol.ConnectionType, addr *net.TCPAddr) {
+	s.changeConnections(connectionType, addr, -1)
+}
+
+func (s *statsStatsd) changeConnections(connectionType newprotocol.ConnectionType, addr *net.TCPAddr, value int) {
+	var labels [3]string
+
+	labels[0] = "connections"
+	switch connectionType {
+	case newprotocol.ConnectionTypeAbridged:
+		labels[1] = "abridged"
+	case newprotocol.ConnectionTypeSecure:
+		labels[1] = "secured"
+	default:
+		labels[1] = "intermediate"
+	}
+
+	labels[2] = "ipv4"
+	if addr.IP.To4() == nil {
+		labels[2] = "ipv6"
+	}
+
+	s.client.Count(strings.Join(labels[:], "."), value)
+}
+
+func (s *statsStatsd) Crash() {
+	s.client.Increment("crashes")
+}
+
+func (s *statsStatsd) AntiReplayDetected() {
+	s.client.Increment("anti_replays")
+}
+
+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),
+	}
+
+	if len(newconfig.C.StatsdStats.Tags) > 0 {
+		tags := make([]string, len(newconfig.C.StatsdStats.Tags)*2)
+		for k, v := range newconfig.C.StatsdStats.Tags {
+			tags = append(tags, k, v)
+		}
+		options = append(options, statsd.Tags(tags...))
+	}
+
+	client, err := statsd.New(options...)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot initialize a client")
+	}
+
+	return &statsStatsd{
+		client: client,
+	}, nil
+}

newwrappers/interfaces.go

+package newwrappers
+
+import (
+	"io"
+	"net"
+	"time"
+
+	"go.uber.org/zap"
+)
+
+type Packet []byte
+
+// Wrap is a base interface for all wrappers in this package.
+type Wrap interface {
+	Conn() net.Conn
+	Logger() *zap.SugaredLogger
+	LocalAddr() *net.TCPAddr
+	RemoteAddr() *net.TCPAddr
+}
+
+type BaseReaderWithTimeout interface {
+	ReadTimeout([]byte, time.Duration) (int, error)
+}
+
+type BaseWriterWithTimeout interface {
+	WriteTimeout([]byte, time.Duration) (int, error)
+}
+
+type BasePacketReader interface {
+	Read() (Packet, error)
+}
+
+type BasePacketWriter interface {
+	Write(Packet) error
+}
+
+type StreamReader interface {
+	Wrap
+	io.Reader
+	BaseReaderWithTimeout
+}
+
+type StreamWriter interface {
+	Wrap
+	io.Writer
+	BaseWriterWithTimeout
+}
+
+type StreamCloser interface {
+	Wrap
+	io.Closer
+}
+
+type StreamReadCloser interface {
+	Wrap
+	io.ReadCloser
+	BaseReaderWithTimeout
+}
+
+type StreamWriteCloser interface {
+	Wrap
+	io.WriteCloser
+	BaseWriterWithTimeout
+}
+
+type StreamReadWriter interface {
+	Wrap
+	io.ReadWriter
+	BaseReaderWithTimeout
+}
+
+type StreamReadWriteCloser interface {
+	Wrap
+	io.ReadWriteCloser
+	BaseReaderWithTimeout
+	BaseWriterWithTimeout
+}
+
+type PacketReader interface {
+	Wrap
+	BasePacketReader
+}
+
+type PacketWriter interface {
+	Wrap
+	BasePacketWriter
+}
+
+type PacketCloser interface {
+	Wrap
+	io.Closer
+}
+
+type PacketReadCloser interface {
+	Wrap
+	BasePacketReader
+	io.Closer
+}
+
+type PacketWriteCloser interface {
+	Wrap
+	BasePacketWriter
+	io.Closer
+}
+
+type PacketReadWriter interface {
+	Wrap
+	BasePacketWriter
+	BasePacketReader
+}
+
+type PacketReadWriteCloser interface {
+	Wrap
+	BasePacketWriter
+	BasePacketReader
+	io.Closer
+}

newwrappers/wrapper_conn.go

+package newwrappers
+
+import (
+	"context"
+	"crypto/rand"
+	"encoding/hex"
+	"net"
+	"time"
+
+	"github.com/juju/errors"
+	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/newconfig"
+)
+
+const ConnIDLength = 8
+
+type ConnID [ConnIDLength]byte
+
+func (c ConnID) String() string {
+	return hex.EncodeToString(c[:])
+}
+
+type connPurpose uint8
+
+const (
+	connPurposeClient connPurpose = 1 << iota
+	connPurposeTelegram
+)
+
+const (
+	connTimeoutRead  = 2 * time.Minute
+	connTimeoutWrite = 2 * time.Minute
+)
+
+type wrapperConn struct {
+	parent     net.Conn
+	ctx        context.Context
+	cancel     context.CancelFunc
+	connID     ConnID
+	logger     *zap.SugaredLogger
+	localAddr  *net.TCPAddr
+	remoteAddr *net.TCPAddr
+}
+
+func (w *wrapperConn) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
+	select {
+	case <-w.ctx.Done():
+		w.Close()
+		return 0, errors.Annotate(w.ctx.Err(), "Cannot write because context was closed")
+
+	default:
+		if err := w.parent.SetWriteDeadline(time.Now().Add(timeout)); err != nil {
+			w.Close() // nolint: gosec
+			return 0, errors.Annotate(err, "Cannot set write deadline to the socket")
+		}
+
+		n, err := w.parent.Write(p)
+		w.logger.Debugw("Write to stream", "bytes", n, "error", err)
+		if err != nil {
+			w.Close() // nolint: gosec
+		}
+
+		return n, err
+	}
+}
+
+func (w *wrapperConn) Write(p []byte) (int, error) {
+	return w.WriteTimeout(p, connTimeoutWrite)
+}
+
+func (w *wrapperConn) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
+	select {
+	case <-w.ctx.Done():
+		w.Close()
+		return 0, errors.Annotate(w.ctx.Err(), "Cannot read because context was closed")
+
+	default:
+		if err := w.parent.SetReadDeadline(time.Now().Add(timeout)); err != nil {
+			w.Close()
+			return 0, errors.Annotate(err, "Cannot set read deadline to the socket")
+		}
+
+		n, err := w.parent.Read(p)
+		w.logger.Debugw("Read from stream", "bytes", n, "error", err)
+		if err != nil {
+			w.Close()
+		}
+
+		return n, err
+	}
+}
+
+func (w *wrapperConn) Read(p []byte) (int, error) {
+	return w.ReadTimeout(p, connTimeoutRead)
+}
+
+func (w *wrapperConn) Close() error {
+	w.logger.Debugw("Close connection")
+	w.cancel()
+
+	return w.parent.Close()
+}
+
+func (w *wrapperConn) Conn() net.Conn {
+	return w.parent
+}
+
+func (w *wrapperConn) Logger() *zap.SugaredLogger {
+	return w.logger
+}
+
+func (w *wrapperConn) LocalAddr() *net.TCPAddr {
+	return w.localAddr
+}
+
+func (w *wrapperConn) RemoteAddr() *net.TCPAddr {
+	return w.remoteAddr
+}
+
+func newConn(ctx context.Context,
+	cancel context.CancelFunc,
+	parent net.Conn,
+	connID ConnID,
+	purpose connPurpose) StreamReadWriteCloser {
+	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
+		}
+	} else if newconfig.C.PublicIPv6Addr.IP != nil {
+		localAddr.IP = newconfig.C.PublicIPv6Addr.IP
+	}
+
+	logger := zap.S().With(
+		"local_address", localAddr,
+		"remote_address", parent.RemoteAddr(),
+	).Named("conn")
+
+	if purpose == connPurposeClient {
+		logger = logger.With("connection_id", connID.String())
+	}
+
+	return &wrapperConn{
+		parent:     parent,
+		ctx:        ctx,
+		cancel:     cancel,
+		connID:     connID,
+		logger:     logger,
+		remoteAddr: parent.RemoteAddr().(*net.TCPAddr),
+		localAddr:  &localAddr,
+	}
+}
+
+func NewClientConn(ctx context.Context,
+	cancel context.CancelFunc,
+	parent net.Conn,
+	connID ConnID) StreamReadWriteCloser {
+	return newConn(ctx, cancel, parent, connID, connPurposeClient)
+}
+
+func NewTelegramConn(ctx context.Context,
+	cancel context.CancelFunc,
+	parent net.Conn,
+	connID ConnID) StreamReadWriteCloser {
+	return newConn(ctx, cancel, parent, connID, connPurposeTelegram)
+}
+
+func NewConnID() ConnID {
+	var id ConnID
+
+	if _, err := rand.Read(id[:]); err != nil {
+		panic(err)
+	}
+
+	return id
+}

newwrappers/wrapper_obfuscated2.go

+package newwrappers
+
+import (
+	"crypto/cipher"
+	"net"
+	"time"
+
+	"github.com/juju/errors"
+	"go.uber.org/zap"
+)
+
+type wrapperObfuscated2 struct {
+	encryptor cipher.Stream
+	decryptor cipher.Stream
+	parent    StreamReadWriteCloser
+}
+
+func (w *wrapperObfuscated2) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
+	n, err := w.parent.ReadTimeout(p, timeout)
+	if err != nil {
+		return 0, errors.Annotate(err, "Cannot read stream ciphered data")
+	}
+	w.decryptor.XORKeyStream(p, p[:n])
+
+	return n, nil
+}
+
+func (w *wrapperObfuscated2) Read(p []byte) (int, error) {
+	n, err := w.parent.Read(p)
+	if err != nil {
+		return 0, errors.Annotate(err, "Cannot read stream ciphered data")
+	}
+	w.decryptor.XORKeyStream(p, p[:n])
+
+	return n, nil
+}
+
+func (w *wrapperObfuscated2) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
+	buf := make([]byte, len(p))
+	copy(buf, p)
+	w.encryptor.XORKeyStream(buf, buf)
+
+	return w.parent.WriteTimeout(buf, timeout)
+}
+
+func (w *wrapperObfuscated2) Write(p []byte) (int, error) {
+	buf := make([]byte, len(p))
+	copy(buf, p)
+	w.encryptor.XORKeyStream(buf, buf)
+
+	return w.parent.Write(buf)
+}
+
+func (w *wrapperObfuscated2) Conn() net.Conn {
+	return w.parent.Conn()
+}
+
+func (w *wrapperObfuscated2) Logger() *zap.SugaredLogger {
+	return w.parent.Logger().Named("obfuscated2")
+}
+
+func (w *wrapperObfuscated2) LocalAddr() *net.TCPAddr {
+	return w.parent.LocalAddr()
+}
+
+func (w *wrapperObfuscated2) RemoteAddr() *net.TCPAddr {
+	return w.parent.RemoteAddr()
+}
+
+func (w *wrapperObfuscated2) Close() error {
+	return w.parent.Close()
+}
+
+func NewObfuscated2(socket StreamReadWriteCloser, encryptor, decryptor cipher.Stream) StreamReadWriteCloser {
+	return &wrapperObfuscated2{
+		parent:    socket,
+		encryptor: encryptor,
+		decryptor: decryptor,
+	}
+}

newwrappers/wrapper_stats.go

+package newwrappers
+
+import (
+	"net"
+	"time"
+
+	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/newstats"
+)
+
+type wrapperStats struct {
+	parent StreamReadWriteCloser
+}
+
+func (w *wrapperStats) Write(p []byte) (int, error) {
+	n, err := w.parent.Write(p)
+	newstats.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)
+
+	return n, err
+}
+
+func (w *wrapperStats) Read(p []byte) (int, error) {
+	n, err := w.parent.Read(p)
+	newstats.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)
+
+	return n, err
+}
+
+func (w *wrapperStats) Conn() net.Conn {
+	return w.parent.Conn()
+}
+
+func (w *wrapperStats) Logger() *zap.SugaredLogger {
+	return w.parent.Logger().Named("traffic")
+}
+
+func (w *wrapperStats) LocalAddr() *net.TCPAddr {
+	return w.parent.LocalAddr()
+}
+
+func (w *wrapperStats) RemoteAddr() *net.TCPAddr {
+	return w.parent.RemoteAddr()
+}
+
+func (w *wrapperStats) Close() error {
+	return w.parent.Close()
+}
+
+func NewTraffic(parent StreamReadWriteCloser) StreamReadWriteCloser {
+	return &wrapperStats{parent}
+}

ntp/ntp.go

 
 const autoUpdatePeriod = time.Minute
 
-var ntpEndpoints = []string{
+var ntpEndpoints = [...]string{
 	"0.pool.ntp.org",
 	"1.pool.ntp.org",
 	"2.pool.ntp.org",