bug fixes

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 {

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:

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

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()
 

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
 }
 

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
 }
 

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 = &registry{
-		ctx:   ctx,
-		conns: map[string]*ctxChannel{},
-	}
-}

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)
 	}

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",
 	}