diff --git a/cli/proxy.go b/cli/proxy.go index 54a0d62..e2d8b20 100644 --- a/cli/proxy.go +++ b/cli/proxy.go @@ -10,6 +10,7 @@ import ( "github.com/9seconds/mtg/antireplay" "github.com/9seconds/mtg/config" + "github.com/9seconds/mtg/hub" "github.com/9seconds/mtg/ntp" "github.com/9seconds/mtg/obfuscated2" "github.com/9seconds/mtg/proxy" @@ -66,6 +67,7 @@ func Proxy() error { Fatal(err) } telegram.Init() + hub.Init(ctx) proxyListener, err := net.Listen("tcp", config.C.Bind.String()) if err != nil { diff --git a/config/config.go b/config/config.go index 543a7a8..a8bf8fd 100644 --- a/config/config.go +++ b/config/config.go @@ -9,6 +9,7 @@ import ( "net" "time" + "github.com/alecthomas/units" "go.uber.org/zap" statsd "gopkg.in/alexcesaro/statsd.v2" ) @@ -104,8 +105,14 @@ func Init(options ...Opt) error { // nolint: gocyclo, funlen C.Bind = opt.Value.(*net.TCPAddr) case OptionTypePublicIPv4: C.PublicIPv4 = opt.Value.(*net.TCPAddr) + if C.PublicIPv4 == nil { + C.PublicIPv4 = &net.TCPAddr{} + } case OptionTypePublicIPv6: C.PublicIPv6 = opt.Value.(*net.TCPAddr) + if C.PublicIPv6 == nil { + C.PublicIPv6 = &net.TCPAddr{} + } case OptionTypeStatsBind: C.StatsBind = opt.Value.(*net.TCPAddr) case OptionTypeStatsNamespace: @@ -133,9 +140,9 @@ func Init(options ...Opt) error { // nolint: gocyclo, funlen case OptionTypeStatsdTags: C.StatsdTags = opt.Value.(map[string]string) case OptionTypeWriteBufferSize: - C.WriteBuffer = int(opt.Value.(uint32)) + C.WriteBuffer = int(opt.Value.(units.Base2Bytes)) case OptionTypeReadBufferSize: - C.ReadBuffer = int(opt.Value.(uint32)) + C.ReadBuffer = int(opt.Value.(units.Base2Bytes)) case OptionTypeAntiReplayMaxSize: C.AntiReplayMaxSize = opt.Value.(int) case OptionTypeAntiReplayEvictionTime: diff --git a/go.mod b/go.mod index b7f80f5..7e21076 100644 --- a/go.mod +++ b/go.mod @@ -5,7 +5,7 @@ replace github.com/golang/lint => github.com/golang/lint v0.0.0-20190227174305-8 require ( github.com/OneOfOne/xxhash v1.2.5 // indirect github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 // indirect - github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4 // indirect + github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4 github.com/allegro/bigcache v1.2.1 github.com/beevik/ntp v0.2.0 github.com/cespare/xxhash v1.1.0 diff --git a/hub/connection.go b/hub/connection.go index 4d2ba5f..ee6f4ec 100644 --- a/hub/connection.go +++ b/hub/connection.go @@ -109,6 +109,7 @@ func newConnection(req *protocol.TelegramRequest, hub *connectionHub) (*connecti conn: conn, hub: hub, id: rand.Int(), // nolint: gosec + done: make(chan struct{}), } go rv.run() diff --git a/hub/connection_hub.go b/hub/connection_hub.go index f39508d..a1101b0 100644 --- a/hub/connection_hub.go +++ b/hub/connection_hub.go @@ -43,11 +43,16 @@ func (c *connectionHub) run() { } func (c *connectionHub) runGC() { + logger := c.logger.Named("gc") + for key, conn := range c.sockets { switch { case conn.closed(): + logger.Debugw("Delete closed socket", "key", key) delete(c.sockets, key) + case conn.idle(): + logger.Debugw("Delete idle socket", "key", key) conn.shutdown() delete(c.sockets, key) return @@ -56,9 +61,14 @@ func (c *connectionHub) runGC() { } func (c *connectionHub) runConnectionRequest(req *connectionHubRequest) { + logger := c.logger.Named("request").With("connection-id", req.request.ConnID) + for key, conn := range c.sockets { delete(c.sockets, key) if !conn.closed() { + logger.Debugw("Choose connection", + "id", conn.id, + "remote_addr", conn.conn.RemoteAddr()) req.response <- conn close(req.response) return @@ -66,16 +76,23 @@ func (c *connectionHub) runConnectionRequest(req *connectionHubRequest) { } if conn, err := newConnection(req.request, c); err == nil { + logger.Debugw("New connection", + "id", conn.id, + "remote_addr", conn.conn.RemoteAddr()) req.response <- conn } close(req.response) } func (c *connectionHub) runBrokenSocket(id int) { + c.logger.Named("broken-socket").Debugw("Delete broken socket", "id", id) delete(c.sockets, id) } func (c *connectionHub) runReturnConnection(conn *connection) { + c.logger.Named("return-connection").Debugw("Return connection", + "id", conn.id, + "remote_addr", conn.conn.RemoteAddr()) c.sockets[conn.id] = conn } diff --git a/hub/hub.go b/hub/hub.go index 09c76d1..e0545e2 100644 --- a/hub/hub.go +++ b/hub/hub.go @@ -32,8 +32,11 @@ func (h *hub) Write(packet conntypes.Packet, req *protocol.TelegramRequest) erro } if err := conn.write(packet); err != nil { + conn.shutdown() return fmt.Errorf("cannot send packet: %w", err) } + sub.channelReturnConnections <- conn + return nil } diff --git a/hub/registry.go b/hub/registry.go index 4beebe8..5f6e33c 100644 --- a/hub/registry.go +++ b/hub/registry.go @@ -42,10 +42,3 @@ func (r *registry) getChannel(id conntypes.ConnID) (*ctxChannel, bool) { } return nil, false } - -func InitRegistry(ctx context.Context) { - Registry = ®istry{ - ctx: ctx, - conns: map[string]*ctxChannel{}, - } -} diff --git a/stats/stats_prometheus.go b/stats/stats_prometheus.go index a075940..bcb0c91 100644 --- a/stats/stats_prometheus.go +++ b/stats/stats_prometheus.go @@ -77,7 +77,7 @@ func (s *statsPrometheus) changeTelegramConnections(dc conntypes.DC, addr *net.T labels[1] = "ipv6" } - s.connections.WithLabelValues(labels[:]...).Add(increment) + s.telegramConnections.WithLabelValues(labels[:]...).Add(increment) } func (s *statsPrometheus) Crash() { @@ -122,6 +122,9 @@ func newStatsPrometheus(mux *http.ServeMux) (Interface, error) { if err := registry.Register(instance.connections); err != nil { return nil, fmt.Errorf("cannot register metrics for connections: %w", err) } + if err := registry.Register(instance.telegramConnections); err != nil { + return nil, fmt.Errorf("cannot register metrics for telegram connections: %w", err) + } if err := registry.Register(instance.traffic); err != nil { return nil, fmt.Errorf("cannot register metrics for traffic: %w", err) } diff --git a/stats/stats_statsd.go b/stats/stats_statsd.go index d5e5d95..9451ed7 100644 --- a/stats/stats_statsd.go +++ b/stats/stats_statsd.go @@ -63,7 +63,7 @@ func (s *statsStatsd) TelegramDisconnected(dc conntypes.DC, addr *net.TCPAddr) { func (s *statsStatsd) changeTelegramConnections(dc conntypes.DC, addr *net.TCPAddr, value int) { labels := [...]string{ - "telegram", + "telegram_connections", strconv.Itoa(int(dc)), "ipv4", }