main.go

package main

import (
	"math/rand"
	"os"
	"runtime/debug"
	"strings"
	"time"

	"github.com/9seconds/mtg/cli"
	"github.com/9seconds/mtg/config"
	"github.com/9seconds/mtg/utils"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

var version = "dev" // has to be set by ldflags

var (
	app = kingpin.New("mtg", "Simple MTPROTO proxy.")

	generateSecretCommand = app.Command("generate-secret",
		"Generate new secret")
	generateCloakHost = generateSecretCommand.Flag("cloak-host",
		"A host to use for TLS cloaking.").
		Short('c').
		Default("storage.googleapis.com").
		String()
	generateSecretType = generateSecretCommand.Arg("type",
		"A type of secret to generate. Valid options are 'simple', 'secured' and 'tls'").
		Required().
		Enum("simple", "secured", "tls")

	runCommand = app.Command("run",
		"Run new proxy instance")
	runDebug = runCommand.Flag("debug",
		"Run in debug mode.").
		Short('d').
		Envar("MTG_DEBUG").
		Bool()
	runVerbose = runCommand.Flag("verbose",
		"Run in verbose mode.").
		Short('v').
		Envar("MTG_VERBOSE").
		Bool()
	runPreferIP = runCommand.Flag("prefer-ip",
		"Prefer this IP protocol if possible. Valid options are 'ipv4' and 'ipv6'").
		Envar("MTG_PREFER_DIRECT_IP").
		Default("ipv6").
		Enum("ipv4", "ipv6")
	runBind = runCommand.Flag("bind",
		"Host:Port to bind proxy to.").
		Short('b').
		Envar("MTG_BIND").
		Default("0.0.0.0:3128").
		TCP()
	runPublicIPv4 = runCommand.Flag("public-ipv4",
		"Which IPv4 host:port to use.").
		Short('4').
		Envar("MTG_IPV4").
		TCP()
	runPublicIPv6 = runCommand.Flag("public-ipv6",
		"Which IPv6 host:port to use.").
		Short('6').
		Envar("MTG_IPV6").
		TCP()
	runStatsBind = runCommand.Flag("stats-bind",
		"Which Host:Port to bind stats server to.").
		Short('t').
		Envar("MTG_STATS_BIND").
		Default("127.0.0.1:3129").
		TCP()
	runStatsNamespace = runCommand.Flag("stats-namespace",
		"Which namespace to use for Prometheus.").
		Envar("MTG_STATS_NAMESPACE").
		Default("mtg").
		String()
	runStatsdAddress = runCommand.Flag("statsd-addr",
		"Host:port of statsd server").
		Envar("MTG_STATSD_ADDR").
		TCP()
	runStatsdTagsFormat = runCommand.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").
		Default("influxdb").
		Enum("datadog", "influxdb")
	runStatsdTags = runCommand.Flag("statsd-tags",
		"Tags to use for working with statsd (specified as 'key=value').").
		Envar("MTG_STATSD_TAGS").
		StringMap()
	runWriteBufferSize = runCommand.Flag("write-buffer",
		"Write buffer size. You can think about it as a buffer from client to Telegram.").
		Short('w').
		Envar("MTG_BUFFER_WRITE").
		Default("32KB").
		Bytes()
	runReadBufferSize = runCommand.Flag("read-buffer",
		"Read buffer size. You can think about it as a buffer from Telegram to client.").
		Short('r').
		Envar("MTG_BUFFER_READ").
		Default("32KB").
		Bytes()
	runTLSCloakPort = runCommand.Flag("cloak-port",
		"Port which should be used for host cloaking.").
		Envar("MTG_CLOAK_PORT").
		Default("443").
		Uint16()
	runAntiReplayMaxSize = runCommand.Flag("anti-replay-max-size",
		"Max size of antireplay cache.").
		Envar("MTG_ANTIREPLAY_MAXSIZE").
		Default("128MB").
		Bytes()
	runMultiplexPerConnection = runCommand.Flag("multiplex-per-connection",
		"How many clients can share a single connection to Telegram.").
		Envar("MTG_MULTIPLEX_PERCONNECTION").
		Default("50").
		Uint()
	runNTPServers = runCommand.Flag("ntp-server",
		"A list of NTP servers to use.").
		Envar("MTG_NTP_SERVERS").
		Default("0.pool.ntp.org", "1.pool.ntp.org", "2.pool.ntp.org", "3.pool.ntp.org").
		Strings()
	runSecret = runCommand.Arg("secret", "Secret of this proxy.").Required().HexBytes()
	runAdtag  = runCommand.Arg("adtag", "ADTag of the proxy.").HexBytes()
)

func main() {
	rand.Seed(time.Now().UTC().UnixNano())
	app.Version(getVersion())
	app.HelpFlag.Short('h')

	if err := utils.SetLimits(); err != nil {
		cli.Fatal(err)
	}

	switch kingpin.MustParse(app.Parse(os.Args[1:])) {
	case generateSecretCommand.FullCommand():
		cli.Generate(*generateSecretType, *generateCloakHost)
	case runCommand.FullCommand():
		err := config.Init(
			config.Opt{Option: config.OptionTypeDebug, Value: *runDebug},
			config.Opt{Option: config.OptionTypeVerbose, Value: *runVerbose},
			config.Opt{Option: config.OptionTypePreferIP, Value: *runPreferIP},
			config.Opt{Option: config.OptionTypeBind, Value: *runBind},
			config.Opt{Option: config.OptionTypePublicIPv4, Value: *runPublicIPv4},
			config.Opt{Option: config.OptionTypePublicIPv6, Value: *runPublicIPv6},
			config.Opt{Option: config.OptionTypeStatsBind, Value: *runStatsBind},
			config.Opt{Option: config.OptionTypeStatsNamespace, Value: *runStatsNamespace},
			config.Opt{Option: config.OptionTypeStatsdAddress, Value: *runStatsdAddress},
			config.Opt{Option: config.OptionTypeStatsdTagsFormat, Value: *runStatsdTagsFormat},
			config.Opt{Option: config.OptionTypeStatsdTags, Value: *runStatsdTags},
			config.Opt{Option: config.OptionTypeWriteBufferSize, Value: *runWriteBufferSize},
			config.Opt{Option: config.OptionTypeReadBufferSize, Value: *runReadBufferSize},
			config.Opt{Option: config.OptionTypeCloakPort, Value: *runTLSCloakPort},
			config.Opt{Option: config.OptionTypeAntiReplayMaxSize, Value: *runAntiReplayMaxSize},
			config.Opt{Option: config.OptionTypeMultiplexPerConnection, Value: *runMultiplexPerConnection},
			config.Opt{Option: config.OptionTypeNTPServers, Value: *runNTPServers},
			config.Opt{Option: config.OptionTypeSecret, Value: *runSecret},
			config.Opt{Option: config.OptionTypeAdtag, Value: *runAdtag},
		)
		if err != nil {
			cli.Fatal(err)
		}

		if err := cli.Proxy(); err != nil {
			cli.Fatal(err)
		}
	}
}

func getVersion() string {
	if version != "dev" {
		return version
	}

	info, ok := debug.ReadBuildInfo()
	if !ok {
		return version
	}

	builder := strings.Builder{}
	builder.WriteString(info.Main.Version)

	if info.Main.Sum != "" {
		builder.WriteString(" (checksum: ")
		builder.WriteString(info.Main.Sum)
		builder.WriteRune(')')
	}

	return builder.String()
}