More reasonable shutdowns

cli/proxy.go

@@ -15,9 +15,12 @@ import (
 	"github.com/9seconds/mtg/proxy"
 	"github.com/9seconds/mtg/stats"
 	"github.com/9seconds/mtg/telegram"
+	"github.com/9seconds/mtg/utils"
 )
 
 func Proxy() error {
+	ctx := utils.GetSignalContext()
+
 	atom := zap.NewAtomicLevel()
 	switch {
 	case config.C.Debug:
@@ -37,7 +40,7 @@ func Proxy() error {
 	zap.ReplaceGlobals(logger)
 	defer logger.Sync() // nolint: errcheck
 
-	if err := config.InitPublicAddress(); err != nil {
+	if err := config.InitPublicAddress(ctx); err != nil {
 		Fatal(err.Error())
 	}
 	zap.S().Debugw("Configuration", "config", config.C)
@@ -61,16 +64,21 @@ func Proxy() error {
 	if err := antireplay.Init(); err != nil {
 		Fatal(err.Error())
 	}
-	if err := stats.Init(); err != nil {
+	if err := stats.Init(ctx); err != nil {
 		Fatal(err.Error())
 	}
 	proxyListener, err := net.Listen("tcp", config.C.ListenAddr.String())
 	if err != nil {
 		Fatal(err.Error())
 	}
+	go func() {
+		<-ctx.Done()
+		proxyListener.Close()
+	}()
 
 	app := &proxy.Proxy{
-		Logger: zap.S().Named("proxy"),
+		Logger:  zap.S().Named("proxy"),
+		Context: ctx,
 	}
 	if len(config.C.AdTag) == 0 {
 		app.TelegramProtocolMaker = obfuscated2.MakeTelegramProtocol

config/config.go

@@ -2,6 +2,7 @@ package config
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"net"
 	"strconv"
@@ -140,14 +141,14 @@ func (c Config) String() string {
 	return string(data)
 }
 
-type ConfigOpt struct {
+type Opt struct {
 	Option OptionType
 	Value  interface{}
 }
 
 var C = Config{}
 
-func Init(options ...ConfigOpt) error { // nolint: gocyclo
+func Init(options ...Opt) error { // nolint: gocyclo
 	for _, opt := range options {
 		switch opt.Option {
 		case OptionTypeDebug:
@@ -222,7 +223,7 @@ func Init(options ...ConfigOpt) error { // nolint: gocyclo
 	return nil
 }
 
-func InitPublicAddress() error {
+func InitPublicAddress(ctx context.Context) error {
 	if C.PublicIPv4Addr.Port == 0 {
 		C.PublicIPv4Addr.Port = C.ListenAddr.Port
 	}
@@ -232,7 +233,7 @@ func InitPublicAddress() error {
 
 	foundAddress := C.PublicIPv4Addr.IP != nil || C.PublicIPv6Addr.IP != nil
 	if C.PublicIPv4Addr.IP == nil {
-		ip, err := getGlobalIPv4()
+		ip, err := getGlobalIPv4(ctx)
 		if err != nil {
 			zap.S().Warnw("Cannot resolve public address", "error", err)
 		} else {
@@ -241,7 +242,7 @@ func InitPublicAddress() error {
 		}
 	}
 	if C.PublicIPv6Addr.IP == nil {
-		ip, err := getGlobalIPv6()
+		ip, err := getGlobalIPv6(ctx)
 		if err != nil {
 			zap.S().Warnw("Cannot resolve public address", "error", err)
 		} else {

config/global_ips.go

@@ -17,23 +17,23 @@ const (
 	ifconfigTimeout = 10 * time.Second
 )
 
-func getGlobalIPv4() (net.IP, error) {
-	ip, err := fetchIP("tcp4")
+func getGlobalIPv4(ctx context.Context) (net.IP, error) {
+	ip, err := fetchIP(ctx, "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) {
-	ip, err := fetchIP("tcp6")
+func getGlobalIPv6(ctx context.Context) (net.IP, error) {
+	ip, err := fetchIP(ctx, "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) {
+func fetchIP(ctx context.Context, network string) (net.IP, error) {
 	dialer := &net.Dialer{FallbackDelay: -1}
 	client := &http.Client{
 		Jar:     nil,
@@ -45,7 +45,12 @@ func fetchIP(network string) (net.IP, error) {
 		},
 	}
 
-	resp, err := client.Get(ifconfigAddress)
+	req, err := http.NewRequest("GET", ifconfigAddress, nil)
+	if err != nil {
+		return nil, errors.Annotate(err, "Cannot create a request")
+	}
+
+	resp, err := client.Do(req.WithContext(ctx))
 	if err != nil {
 		if resp != nil {
 			io.Copy(ioutil.Discard, resp.Body) // nolint: errcheck

main.go

@@ -152,29 +152,29 @@ func main() {
 
 	case proxyCommand.FullCommand():
 		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},
+			config.Opt{Option: config.OptionTypeDebug, Value: *proxyDebug},
+			config.Opt{Option: config.OptionTypeVerbose, Value: *proxyVerbose},
+			config.Opt{Option: config.OptionTypeBindIP, Value: *proxyBindIP},
+			config.Opt{Option: config.OptionTypeBindPort, Value: *proxyBindPort},
+			config.Opt{Option: config.OptionTypePublicIPv4, Value: *proxyPublicIPv4},
+			config.Opt{Option: config.OptionTypePublicIPv4Port, Value: *proxyPublicIPv4Port},
+			config.Opt{Option: config.OptionTypePublicIPv6, Value: *proxyPublicIPv6},
+			config.Opt{Option: config.OptionTypePublicIPv6Port, Value: *proxyPublicIPv6Port},
+			config.Opt{Option: config.OptionTypeStatsIP, Value: *proxyStatsIP},
+			config.Opt{Option: config.OptionTypeStatsPort, Value: *proxyStatsPort},
+			config.Opt{Option: config.OptionTypeStatsdIP, Value: *proxyStatsdIP},
+			config.Opt{Option: config.OptionTypeStatsdPort, Value: *proxyStatsdPort},
+			config.Opt{Option: config.OptionTypeStatsdNetwork, Value: *proxyStatsdNetwork},
+			config.Opt{Option: config.OptionTypeStatsdPrefix, Value: *proxyStatsdPrefix},
+			config.Opt{Option: config.OptionTypeStatsdTagsFormat, Value: *proxyStatsdTagsFormat},
+			config.Opt{Option: config.OptionTypeStatsdTags, Value: *proxyStatsdTags},
+			config.Opt{Option: config.OptionTypePrometheusPrefix, Value: *proxyPrometheusPrefix},
+			config.Opt{Option: config.OptionTypeWriteBufferSize, Value: *proxyWriteBufferSize},
+			config.Opt{Option: config.OptionTypeReadBufferSize, Value: *proxyReadBufferSize},
+			config.Opt{Option: config.OptionTypeAntiReplayMaxSize, Value: *proxyAntiReplayMaxSize},
+			config.Opt{Option: config.OptionTypeAntiReplayEvictionTime, Value: *proxyAntiReplayEvictionTime},
+			config.Opt{Option: config.OptionTypeSecret, Value: *proxySecret},
+			config.Opt{Option: config.OptionTypeAdtag, Value: *proxyAdtag},
 		)
 		if err != nil {
 			cli.Fatal(err.Error())

proxy/proxy.go

@@ -20,17 +20,25 @@ const directPipeBufferSize = 1024 * 1024
 
 type Proxy struct {
 	Logger                *zap.SugaredLogger
+	Context               context.Context
 	ClientProtocolMaker   protocol.ClientProtocolMaker
 	TelegramProtocolMaker protocol.TelegramProtocolMaker
 	TelegramDialer        telegram.Telegram
 }
 
 func (p *Proxy) Serve(listener net.Listener) {
+	doneChan := p.Context.Done()
+
 	for {
 		conn, err := listener.Accept()
 		if err != nil {
-			p.Logger.Errorw("Cannot allocate incoming connection", "error", err)
-			continue
+			select {
+			case <-doneChan:
+				return
+			default:
+				p.Logger.Errorw("Cannot allocate incoming connection", "error", err)
+				continue
+			}
 		}
 		go p.accept(conn)
 	}
@@ -53,7 +61,7 @@ func (p *Proxy) accept(conn net.Conn) {
 		return
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(p.Context)
 	defer cancel()
 
 	wrappedConn := wrappers.NewClientConn(ctx, cancel, conn, connID)

stats/stats.go

@@ -1,6 +1,7 @@
 package stats
 
 import (
+	"context"
 	"net"
 	"net/http"
 
@@ -59,7 +60,7 @@ func (m multiStats) AntiReplayDetected() {
 
 var S Stats
 
-func Init() error {
+func Init(ctx context.Context) error {
 	mux := http.NewServeMux()
 
 	instanceJSON := newStatsJSON(mux)
@@ -86,6 +87,10 @@ func Init() error {
 		Handler: mux,
 	}
 	go srv.Serve(listener) // nolint: errcheck
+	go func() {
+		<-ctx.Done()
+		srv.Shutdown(context.Background()) // nolint: errcheck
+	}()
 
 	S = multiStats(stats)
 

utils/signal_context.go

@@ -0,0 +1,24 @@
+// +build !windows
+
+package utils
+
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+func GetSignalContext() context.Context {
+	ctx, cancel := context.WithCancel(context.Background())
+	sigChan := make(chan os.Signal, 1)
+
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		for range sigChan {
+			cancel()
+		}
+	}()
+
+	return ctx
+}

utils/signal_context_windows.go

@@ -0,0 +1,23 @@
+// +build windows
+
+package utils
+
+import (
+	"context"
+	"os"
+	"os/signal"
+)
+
+func GetSignalContext() context.Context {
+	ctx, cancel := context.WithCancel(context.Background())
+	sigChan := make(chan os.Signal, 1)
+
+	signal.Notify(sigChan, os.Interrupt)
+	go func() {
+		for range sigChan {
+			cancel()
+		}
+	}()
+
+	return ctx
+}