ScuroNeko/mtg
1
0
Fork 0
mirror of https://github.com/ScuroNeko/mtg.git synced 2026-08-31 14:24:01 +03:00
Code Issues Packages Projects Releases Wiki Activity

More reasonable shutdowns

Browse Source
This commit is contained in:
9seconds
2019-09-04 10:54:38 +03:00
parent 299a34252e
commit 701f62ed4c
8 changed files with 115 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cli/proxy.go
+11 -3
View File
@@ -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
+6 -5
View File
@@ -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
+11 -6
View File
@@ -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
+23 -23
View File
@@ -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
+11 -3
View File
@@ -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
+6 -1
View File
@@ -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
+24
View File
@@ -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
+23
View File
@@ -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
}