diff --git a/cli/proxy.go b/cli/proxy.go index bccddfc..a08f71c 100644 --- a/cli/proxy.go +++ b/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 diff --git a/config/config.go b/config/config.go index d9b9adb..2aec548 100644 --- a/config/config.go +++ b/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 { diff --git a/config/global_ips.go b/config/global_ips.go index f2e72de..719000a 100644 --- a/config/global_ips.go +++ b/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 diff --git a/main.go b/main.go index 4fc07d0..3abdf20 100644 --- a/main.go +++ b/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()) diff --git a/proxy/proxy.go b/proxy/proxy.go index 96108b6..7ead6d9 100644 --- a/proxy/proxy.go +++ b/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) diff --git a/stats/stats.go b/stats/stats.go index 5793552..6412132 100644 --- a/stats/stats.go +++ b/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) diff --git a/utils/signal_context.go b/utils/signal_context.go new file mode 100644 index 0000000..40e2b6f --- /dev/null +++ b/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 +} diff --git a/utils/signal_context_windows.go b/utils/signal_context_windows.go new file mode 100644 index 0000000..d18ce71 --- /dev/null +++ b/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 +}