Add new stats metric

antireplay/init.go

@@ -16,9 +16,9 @@ func Init() {
 	initOnce.Do(func() {
 		c, err := bigcache.NewBigCache(bigcache.Config{
 			Shards:           1024,
-			LifeWindow:       config.C.AntiReplay.EvictionTime,
+			LifeWindow:       config.C.AntiReplayEvictionTime,
 			Hasher:           hasher{},
-			HardMaxCacheSize: config.C.AntiReplay.MaxSize,
+			HardMaxCacheSize: config.C.AntiReplayMaxSize,
 		})
 		if err != nil {
 			panic(err)

cli/proxy.go

@@ -43,7 +43,7 @@ func Proxy() error {
 	if err := config.InitPublicAddress(ctx); err != nil {
 		Fatal(err)
 	}
-	zap.S().Debugw("Configuration", "config", config.C.Printable())
+	zap.S().Debugw("Configuration", "config", config.Printable())
 
 	if len(config.C.AdTag) > 0 {
 		zap.S().Infow("Use middle proxy connection to Telegram")
@@ -67,7 +67,7 @@ func Proxy() error {
 	}
 	telegram.MiddleInit()
 
-	proxyListener, err := net.Listen("tcp", config.C.ListenAddr.String())
+	proxyListener, err := net.Listen("tcp", config.C.Bind.String())
 	if err != nil {
 		Fatal(err)
 	}

config/config.go

@@ -7,7 +7,6 @@ import (
 	"errors"
 	"fmt"
 	"net"
-	"strconv"
 	"time"
 
 	"go.uber.org/zap"
@@ -40,22 +39,16 @@ const (
 	OptionTypeDebug OptionType = iota
 	OptionTypeVerbose
 
-	OptionTypeBindIP
+	OptionTypeBind
-	OptionTypeBindPort
 	OptionTypePublicIPv4
-	OptionTypePublicIPv4Port
 	OptionTypePublicIPv6
-	OptionTypePublicIPv6Port
-	OptionTypeStatsIP
-	OptionTypeStatsPort
 
-	OptionTypeStatsdIP
+	OptionTypeStatsBind
-	OptionTypeStatsdPort
+	OptionTypeStatsNamespace
+	OptionTypeStatsdAddress
 	OptionTypeStatsdNetwork
-	OptionTypeStatsdPrefix
 	OptionTypeStatsdTagsFormat
 	OptionTypeStatsdTags
-	OptionTypePrometheusPrefix
 
 	OptionTypeWriteBufferSize
 	OptionTypeReadBufferSize
@@ -67,93 +60,30 @@ const (
 	OptionTypeAdtag
 )
 
-type BufferSize struct {
-	Read  int `json:"read"`
-	Write int `json:"write"`
-}
-
-type AntiReplay struct {
-	MaxSize      int           `json:"max_size"`
-	EvictionTime time.Duration `json:"duration"`
-}
-
-type Stats struct {
-	Prefix string `json:"prefix"`
-}
-
-type StatsdStats struct {
-	Stats
-
-	Addr       Addr              `json:"addr"`
-	Tags       map[string]string `json:"tags"`
-	TagsFormat statsd.TagFormat  `json:"format"`
-}
-
-type PrometheusStats struct {
-	Stats
-}
-
-type Addr struct {
-	IP   net.IP `json:"ip"`
-	Port int    `json:"port"`
-	net  string
-}
-
-func (a Addr) Network() string {
-	if a.net == "" {
-		return "tcp"
-	}
-	return a.net
-}
-
-func (a Addr) String() string {
-	return net.JoinHostPort(a.IP.String(), strconv.Itoa(a.Port))
-}
-
-func (a Addr) MarshalJSON() ([]byte, error) {
-	data := map[string]string{
-		"network": a.Network(),
-		"addr":    a.String(),
-	}
-	return json.Marshal(data)
-}
-
 type Config struct {
-	BufferSize BufferSize `json:"buffer_size"`
+	Bind       *net.TCPAddr `json:"bind"`
-	AntiReplay AntiReplay `json:"anti_replay"`
+	PublicIPv4 *net.TCPAddr `json:"public_ipv4"`
+	PublicIPv6 *net.TCPAddr `json:"public_ipv6"`
+	StatsBind  *net.TCPAddr `json:"stats_bind"`
+	StatsdAddr *net.TCPAddr `json:"stats_addr"`
 
-	ListenAddr     Addr `json:"listen_addr"`
+	StatsNamespace string            `json:"stats_namespace"`
-	PublicIPv4Addr Addr `json:"public_ipv4_addr"`
+	StatsdNetwork  string            `json:"statsd_network"`
-	PublicIPv6Addr Addr `json:"public_ipv6_addr"`
+	StatsdTags     map[string]string `json:"statsd_tags"`
-	StatsAddr      Addr `json:"stats_addr"`
 
-	StatsdStats     StatsdStats     `json:"stats_statsd"`
+	WriteBuffer int `json:"write_buffer"`
-	PrometheusStats PrometheusStats `json:"stats_prometheus"`
+	ReadBuffer  int `json:"read_buffer"`
 
-	Debug      bool       `json:"debug"`
+	AntiReplayMaxSize      int           `json:"anti_replay_max_size"`
-	Verbose    bool       `json:"verbose"`
+	AntiReplayEvictionTime time.Duration `json:"anti_replay_eviction_time"`
-	SecretMode SecretMode `json:"secret_mode"`
-	Secret     []byte     `json:"secret"`
-	AdTag      []byte     `json:"adtag"`
-}
 
-func (c Config) Printable() interface{} {
+	Debug            bool             `json:"debug"`
-	data, err := json.Marshal(c)
+	Verbose          bool             `json:"verbose"`
-	if err != nil {
+	StatsdTagsFormat statsd.TagFormat `json:"statsd_tags_format"`
-		panic(err)
+	SecretMode       SecretMode       `json:"secret_mode"`
-	}
 
-	rv := map[string]interface{}{}
+	Secret []byte `json:"secret"`
-	if err := json.Unmarshal(data, &rv); err != nil {
+	AdTag  []byte `json:"adtag"`
-		panic(err)
-	}
-
-	return rv
-}
-
-func (c Config) String() string {
-	data, _ := json.Marshal(c)
-	return string(data)
 }
 
 type Opt struct {
@@ -163,59 +93,53 @@ type Opt struct {
 
 var C = Config{}
 
-func Init(options ...Opt) error { // nolint: gocyclo
+func Init(options ...Opt) error { // nolint: gocyclo, funlen
 	for _, opt := range options {
 		switch opt.Option {
 		case OptionTypeDebug:
 			C.Debug = opt.Value.(bool)
 		case OptionTypeVerbose:
 			C.Verbose = opt.Value.(bool)
-		case OptionTypeBindIP:
+		case OptionTypeBind:
-			C.ListenAddr.IP = opt.Value.(net.IP)
+			C.Bind = opt.Value.(*net.TCPAddr)
-		case OptionTypeBindPort:
-			C.ListenAddr.Port = int(opt.Value.(uint16))
 		case OptionTypePublicIPv4:
-			C.PublicIPv4Addr.IP = opt.Value.(net.IP)
+			C.PublicIPv4 = opt.Value.(*net.TCPAddr)
-		case OptionTypePublicIPv4Port:
-			C.PublicIPv4Addr.Port = int(opt.Value.(uint16))
 		case OptionTypePublicIPv6:
-			C.PublicIPv6Addr.IP = opt.Value.(net.IP)
+			C.PublicIPv6 = opt.Value.(*net.TCPAddr)
-		case OptionTypePublicIPv6Port:
+		case OptionTypeStatsBind:
-			C.PublicIPv6Addr.Port = int(opt.Value.(uint16))
+			C.StatsBind = opt.Value.(*net.TCPAddr)
-		case OptionTypeStatsIP:
+		case OptionTypeStatsNamespace:
-			C.StatsAddr.IP = opt.Value.(net.IP)
+			C.StatsNamespace = opt.Value.(string)
-		case OptionTypeStatsPort:
+		case OptionTypeStatsdAddress:
-			C.StatsAddr.Port = int(opt.Value.(uint16))
+			C.StatsdAddr = opt.Value.(*net.TCPAddr)
-		case OptionTypeStatsdIP:
-			C.StatsdStats.Addr.IP = opt.Value.(net.IP)
-		case OptionTypeStatsdPort:
-			C.StatsdStats.Addr.Port = int(opt.Value.(uint16))
 		case OptionTypeStatsdNetwork:
-			C.StatsdStats.Addr.net = opt.Value.(string)
+			value := opt.Value.(string)
-		case OptionTypeStatsdPrefix:
+			switch value {
-			C.StatsdStats.Prefix = opt.Value.(string)
+			case "udp", "tcp":
+				C.StatsdNetwork = value
+			default:
+				return fmt.Errorf("unknown statsd network %v", value)
+			}
 		case OptionTypeStatsdTagsFormat:
 			value := opt.Value.(string)
 			switch value {
 			case "datadog":
-				C.StatsdStats.TagsFormat = statsd.Datadog
+				C.StatsdTagsFormat = statsd.Datadog
 			case "influxdb":
-				C.StatsdStats.TagsFormat = statsd.InfluxDB
+				C.StatsdTagsFormat = statsd.InfluxDB
 			default:
 				return fmt.Errorf("Incorrect statsd tag %s", value)
 			}
 		case OptionTypeStatsdTags:
-			C.StatsdStats.Tags = opt.Value.(map[string]string)
+			C.StatsdTags = opt.Value.(map[string]string)
-		case OptionTypePrometheusPrefix:
-			C.PrometheusStats.Prefix = opt.Value.(string)
 		case OptionTypeWriteBufferSize:
-			C.BufferSize.Write = int(opt.Value.(uint32))
+			C.WriteBuffer = int(opt.Value.(uint32))
 		case OptionTypeReadBufferSize:
-			C.BufferSize.Read = int(opt.Value.(uint32))
+			C.ReadBuffer = int(opt.Value.(uint32))
 		case OptionTypeAntiReplayMaxSize:
-			C.AntiReplay.MaxSize = opt.Value.(int)
+			C.AntiReplayMaxSize = opt.Value.(int)
 		case OptionTypeAntiReplayEvictionTime:
-			C.AntiReplay.EvictionTime = opt.Value.(time.Duration)
+			C.AntiReplayEvictionTime = opt.Value.(time.Duration)
 		case OptionTypeSecret:
 			C.Secret = opt.Value.([]byte)
 		case OptionTypeAdtag:
@@ -239,29 +163,29 @@ func Init(options ...Opt) error { // nolint: gocyclo
 }
 
 func InitPublicAddress(ctx context.Context) error {
-	if C.PublicIPv4Addr.Port == 0 {
+	if C.PublicIPv4.Port == 0 {
-		C.PublicIPv4Addr.Port = C.ListenAddr.Port
+		C.PublicIPv4.Port = C.Bind.Port
 	}
-	if C.PublicIPv6Addr.Port == 0 {
+	if C.PublicIPv6.Port == 0 {
-		C.PublicIPv6Addr.Port = C.ListenAddr.Port
+		C.PublicIPv6.Port = C.Bind.Port
 	}
 
-	foundAddress := C.PublicIPv4Addr.IP != nil || C.PublicIPv6Addr.IP != nil
+	foundAddress := C.PublicIPv4.IP != nil || C.PublicIPv6.IP != nil
-	if C.PublicIPv4Addr.IP == nil {
+	if C.PublicIPv4.IP == nil {
 		ip, err := getGlobalIPv4(ctx)
 		if err != nil {
 			zap.S().Warnw("Cannot resolve public address", "error", err)
 		} else {
-			C.PublicIPv4Addr.IP = ip
+			C.PublicIPv4.IP = ip
 			foundAddress = true
 		}
 	}
-	if C.PublicIPv6Addr.IP == nil {
+	if C.PublicIPv6.IP == nil {
 		ip, err := getGlobalIPv6(ctx)
 		if err != nil {
 			zap.S().Warnw("Cannot resolve public address", "error", err)
 		} else {
-			C.PublicIPv6Addr.IP = ip
+			C.PublicIPv6.IP = ip
 			foundAddress = true
 		}
 	}
@@ -272,3 +196,22 @@ func InitPublicAddress(ctx context.Context) error {
 
 	return nil
 }
+
+func Printable() interface{} {
+	data, err := json.Marshal(C)
+	if err != nil {
+		panic(err)
+	}
+
+	rv := map[string]interface{}{}
+	if err := json.Unmarshal(data, &rv); err != nil {
+		panic(err)
+	}
+
+	rrv, err := json.Marshal(rv)
+	if err != nil {
+		panic(err)
+	}
+
+	return rrv
+}

config/urls.go

@@ -2,6 +2,7 @@ package config
 
 import (
 	"encoding/hex"
+	"fmt"
 	"net/url"
 )
 
@@ -27,14 +28,14 @@ func GetURLs() (urls IPURLs) {
 		secret = "dd" + hex.EncodeToString(C.Secret)
 	}
 
-	urls.IPv4 = makeURLs(&C.PublicIPv4Addr, secret)
+	urls.IPv4 = makeURLs(C.PublicIPv4, secret)
-	urls.IPv6 = makeURLs(&C.PublicIPv6Addr, secret)
+	urls.IPv6 = makeURLs(C.PublicIPv6, secret)
 	urls.BotSecret = secret
 
 	return urls
 }
 
-func makeURLs(addr *Addr, secret string) (urls URLs) {
+func makeURLs(addr fmt.Stringer, secret string) (urls URLs) {
 	values := url.Values{}
 	values.Set("address", addr.String())
 	values.Set("secret", secret)

hub/connection.go

@@ -52,7 +52,7 @@ func (c *connection) write(packet conntypes.Packet) error {
 func (c *connection) shutdown() {
 	c.shutdownOnce.Do(func() {
 		close(c.done)
-			c.hub.channelBrokenSockets <- c.id
+		c.hub.channelBrokenSockets <- c.id
 	})
 }
 

hub/ctx_channel.go

@@ -2,7 +2,6 @@ package hub
 
 import (
 	"context"
-	"errors"
 	"time"
 
 	"github.com/9seconds/mtg/conntypes"

main.go

@@ -36,67 +36,42 @@ var (
 		Short('v').
 		Envar("MTG_VERBOSE").
 		Bool()
-	proxyBindIP = proxyCommand.Flag("bind-ip",
+	proxyBind = proxyCommand.Flag("bind",
-		"Which IP to bind to.").
+		"Host:Port to bind proxy to.").
 		Short('b').
-		Envar("MTG_IP").
+		Envar("MTG_BIND").
-		Default("127.0.0.1").
+		Default("0.0.0.0:3128").
-		IP()
+		TCP()
-	proxyBindPort = proxyCommand.Flag("bind-port",
-		"Which port to bind to.").
-		Short('p').
-		Envar("MTG_PORT").
-		Default("3128").
-		Uint16()
 	proxyPublicIPv4 = proxyCommand.Flag("public-ipv4",
-		"Which IPv4 address is public.").
+		"Which IPv4 host:port to use.").
 		Short('4').
 		Envar("MTG_IPV4").
-		IP()
+		TCP()
-	proxyPublicIPv4Port = proxyCommand.Flag("public-ipv4-port",
-		"Which IPv4 port is public. Default is 'bind-port' value.").
-		Envar("MTG_IPV4_PORT").
-		Uint16()
 	proxyPublicIPv6 = proxyCommand.Flag("public-ipv6",
-		"Which IPv6 address is public.").
+		"Which IPv6 host:port to use.").
 		Short('6').
 		Envar("MTG_IPV6").
-		IP()
+		TCP()
-	proxyPublicIPv6Port = proxyCommand.Flag("public-ipv6-port",
+	proxyStatsBind = proxyCommand.Flag("stats-bind",
-		"Which IPv6 port is public. Default is 'bind-port' value.").
+		"Which Host:Port to bind stats server to.").
-		Envar("MTG_IPV6_PORT").
-		Uint16()
-	proxyStatsIP = proxyCommand.Flag("stats-ip",
-		"Which IP bind stats server to.").
 		Short('t').
-		Envar("MTG_STATS_IP").
+		Envar("MTG_STATS_BIND").
-		Default("127.0.0.1").
+		Default("127.0.0.1:3129").
-		IP()
+		TCP()
-	proxyStatsPort = proxyCommand.Flag("stats-port",
+	proxyStatsNamespace = proxyCommand.Flag("prometheus-namespace",
-		"Which port bind stats to.").
+		"Which namespace to use for Prometheus.").
-		Short('q').
+		Envar("MTG_STATS_NAMESPACE").
-		Envar("MTG_STATS_PORT").
+		Default("mtg").
-		Default("3129").
+		String()
-		Uint16()
+	proxyStatsdAddress = proxyCommand.Flag("statsd-addr",
-	proxyStatsdIP = proxyCommand.Flag("statsd-ip",
+		"Host:port of statsd server").
-		"Which IP should we use for working with statsd.").
+		Envar("MTG_STATSD_ADDR").
-		Envar("MTG_STATSD_IP").
+		TCP()
-		IP()
-	proxyStatsdPort = proxyCommand.Flag("statsd-port",
-		"Which port should we use for working with statsd.").
-		Envar("MTG_STATSD_PORT").
-		Default("8125").
-		Uint16()
 	proxyStatsdNetwork = proxyCommand.Flag("statsd-network",
 		"Which network is used to work with statsd. Only 'tcp' and 'udp' are supported.").
 		Envar("MTG_STATSD_NETWORK").
 		Default("udp").
 		Enum("udp", "tcp")
-	proxyStatsdPrefix = proxyCommand.Flag("statsd-prefix",
-		"Which bucket prefix should we use for sending stats to statsd.").
-		Envar("MTG_STATSD_PREFIX").
-		Default("mtg").
-		String()
 	proxyStatsdTagsFormat = proxyCommand.Flag("statsd-tags-format",
 		"Which tag format should we use to send stats metrics. Valid options are 'datadog' and 'influxdb'.").
 		Envar("MTG_STATSD_TAGS_FORMAT").
@@ -106,23 +81,18 @@ var (
 		"Tags to use for working with statsd (specified as 'key=value').").
 		Envar("MTG_STATSD_TAGS").
 		StringMap()
-	proxyPrometheusPrefix = proxyCommand.Flag("prometheus-prefix",
-		"Which namespace to use to send stats to Prometheus.").
-		Envar("MTG_PROMETHEUS_PREFIX").
-		Default("mtg").
-		String()
 	proxyWriteBufferSize = proxyCommand.Flag("write-buffer",
 		"Write buffer size in bytes. You can think about it as a buffer from client to Telegram.").
 		Short('w').
 		Envar("MTG_BUFFER_WRITE").
-		Default("65536").
+		Default("65536KB").
-		Uint32()
+		Bytes()
 	proxyReadBufferSize = proxyCommand.Flag("read-buffer",
 		"Read buffer size in bytes. You can think about it as a buffer from Telegram to client.").
 		Short('r').
 		Envar("MTG_BUFFER_READ").
-		Default("131072").
+		Default("131072KB").
-		Uint32()
+		Bytes()
 	proxyAntiReplayMaxSize = proxyCommand.Flag("anti-replay-max-size",
 		"Max size of antireplay cache in megabytes.").
 		Envar("MTG_ANTIREPLAY_MAXSIZE").
@@ -154,21 +124,15 @@ func main() {
 		err := config.Init(
 			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.OptionTypeBind, Value: *proxyBind},
-			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.OptionTypeStatsBind, Value: *proxyStatsBind},
-			config.Opt{Option: config.OptionTypeStatsIP, Value: *proxyStatsIP},
+			config.Opt{Option: config.OptionTypeStatsNamespace, Value: *proxyStatsNamespace},
-			config.Opt{Option: config.OptionTypeStatsPort, Value: *proxyStatsPort},
+			config.Opt{Option: config.OptionTypeStatsdAddress, Value: *proxyStatsdAddress},
-			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},

proxy/proxy.go

@@ -47,7 +47,7 @@ func (p *Proxy) accept(conn net.Conn) {
 	defer func() {
 		conn.Close()
 		if err := recover(); err != nil {
-			stats.S.Crash()
+			stats.Stats.Crash()
 			p.Logger.Errorw("Crash of accept handler", "error", err)
 		}
 	}()
@@ -66,7 +66,6 @@ func (p *Proxy) accept(conn net.Conn) {
 	clientConn := wrappers.NewClientConn(conn, connID)
 	clientConn = wrappers.NewCtx(ctx, cancel, clientConn)
 	clientConn = wrappers.NewTimeout(clientConn)
-	clientConn = wrappers.NewTraffic(clientConn)
 	defer clientConn.Close()
 
 	clientProtocol := p.ClientProtocolMaker()
@@ -76,8 +75,8 @@ func (p *Proxy) accept(conn net.Conn) {
 		return
 	}
 
-	stats.S.ClientConnected(clientProtocol.ConnectionType(), clientConn.RemoteAddr())
+	stats.Stats.ClientConnected(clientProtocol.ConnectionType(), clientConn.RemoteAddr())
-	defer stats.S.ClientDisconnected(clientProtocol.ConnectionType(), clientConn.RemoteAddr())
+	defer stats.Stats.ClientDisconnected(clientProtocol.ConnectionType(), clientConn.RemoteAddr())
 	logger.Infow("Client connected", "addr", conn.RemoteAddr())
 
 	req := &protocol.TelegramRequest{

stats/interfaces.go

@@ -0,0 +1,50 @@
+package stats
+
+import (
+	"net"
+
+	"github.com/9seconds/mtg/conntypes"
+)
+
+type IngressTrafficInterface interface {
+	IngressTraffic(int)
+}
+
+type EgressTrafficInterface interface {
+	EgressTraffic(int)
+}
+
+type ClientConnectedInterface interface {
+	ClientConnected(conntypes.ConnectionType, *net.TCPAddr)
+}
+
+type ClientDisconnectedInterface interface {
+	ClientDisconnected(conntypes.ConnectionType, *net.TCPAddr)
+}
+
+type TelegramConnectedInterface interface {
+	TelegramConnected(conntypes.DC, *net.TCPAddr)
+}
+
+type TelegramDisconnectedInterface interface {
+	TelegramDisconnected(conntypes.DC, *net.TCPAddr)
+}
+
+type CrashInterface interface {
+	Crash()
+}
+
+type AntiReplayDetectedInterface interface {
+	AntiReplayDetected()
+}
+
+type Interface interface {
+	IngressTrafficInterface
+	EgressTrafficInterface
+	ClientConnectedInterface
+	ClientDisconnectedInterface
+	TelegramConnectedInterface
+	TelegramDisconnectedInterface
+	CrashInterface
+	AntiReplayDetectedInterface
+}

stats/multi_stats.go

@@ -0,0 +1,57 @@
+package stats
+
+import (
+	"net"
+
+	"github.com/9seconds/mtg/conntypes"
+)
+
+type multiStats []Interface
+
+func (m multiStats) IngressTraffic(traffic int) {
+	for i := range m {
+		go m[i].IngressTraffic(traffic)
+	}
+}
+
+func (m multiStats) EgressTraffic(traffic int) {
+	for i := range m {
+		go m[i].EgressTraffic(traffic)
+	}
+}
+
+func (m multiStats) ClientConnected(connectionType conntypes.ConnectionType, addr *net.TCPAddr) {
+	for i := range m {
+		go m[i].ClientConnected(connectionType, addr)
+	}
+}
+
+func (m multiStats) ClientDisconnected(connectionType conntypes.ConnectionType, addr *net.TCPAddr) {
+	for i := range m {
+		go m[i].ClientDisconnected(connectionType, addr)
+	}
+}
+
+func (m multiStats) TelegramConnected(dc conntypes.DC, addr *net.TCPAddr) {
+	for i := range m {
+		go m[i].TelegramConnected(dc, addr)
+	}
+}
+
+func (m multiStats) TelegramDisconnected(dc conntypes.DC, addr *net.TCPAddr) {
+	for i := range m {
+		go m[i].TelegramDisconnected(dc, addr)
+	}
+}
+
+func (m multiStats) Crash() {
+	for i := range m {
+		go m[i].Crash()
+	}
+}
+
+func (m multiStats) AntiReplayDetected() {
+	for i := range m {
+		go m[i].AntiReplayDetected()
+	}
+}

stats/stats.go

@@ -7,69 +7,20 @@ import (
 	"net/http"
 
 	"github.com/9seconds/mtg/config"
-	"github.com/9seconds/mtg/conntypes"
 )
 
-type Stats interface {
+var Stats Interface
-	IngressTraffic(int)
-	EgressTraffic(int)
-	ClientConnected(conntypes.ConnectionType, *net.TCPAddr)
-	ClientDisconnected(conntypes.ConnectionType, *net.TCPAddr)
-	Crash()
-	AntiReplayDetected()
-}
-
-type multiStats []Stats
-
-func (m multiStats) IngressTraffic(traffic int) {
-	for i := range m {
-		go m[i].IngressTraffic(traffic)
-	}
-}
-
-func (m multiStats) EgressTraffic(traffic int) {
-	for i := range m {
-		go m[i].EgressTraffic(traffic)
-	}
-}
-
-func (m multiStats) ClientConnected(connectionType conntypes.ConnectionType, addr *net.TCPAddr) {
-	for i := range m {
-		go m[i].ClientConnected(connectionType, addr)
-	}
-}
-
-func (m multiStats) ClientDisconnected(connectionType conntypes.ConnectionType, addr *net.TCPAddr) {
-	for i := range m {
-		go m[i].ClientDisconnected(connectionType, addr)
-	}
-}
-
-func (m multiStats) Crash() {
-	for i := range m {
-		go m[i].Crash()
-	}
-}
-
-func (m multiStats) AntiReplayDetected() {
-	for i := range m {
-		go m[i].AntiReplayDetected()
-	}
-}
-
-var S Stats
 
 func Init(ctx context.Context) error {
 	mux := http.NewServeMux()
 
-	instanceJSON := newStatsJSON(mux)
 	instancePrometheus, err := newStatsPrometheus(mux)
 	if err != nil {
 		return fmt.Errorf("cannot initialize prometheus: %w", err)
 	}
 
-	stats := []Stats{instanceJSON, instancePrometheus}
+	stats := []Interface{instancePrometheus}
-	if config.C.StatsdStats.Addr.IP != nil {
+	if config.C.StatsdAddr != nil {
 		instanceStatsd, err := newStatsStatsd()
 		if err != nil {
 			return fmt.Errorf("cannot inialize statsd: %w", err)
@@ -77,7 +28,7 @@ func Init(ctx context.Context) error {
 		stats = append(stats, instanceStatsd)
 	}
 
-	listener, err := net.Listen("tcp", config.C.StatsAddr.String())
+	listener, err := net.Listen("tcp", config.C.StatsBind.String())
 	if err != nil {
 		return fmt.Errorf("cannot initialize stats server: %w", err)
 	}
@@ -91,7 +42,7 @@ func Init(ctx context.Context) error {
 		srv.Shutdown(context.Background()) // nolint: errcheck
 	}()
 
-	S = multiStats(stats)
+	Stats = multiStats(stats)
 
 	return nil
 }

stats/stats_json.go

@@ -1,131 +0,0 @@
-package stats
-
-import (
-	"encoding/json"
-	"net"
-	"net/http"
-	"strconv"
-	"sync/atomic"
-	"time"
-
-	"go.uber.org/zap"
-
-	"github.com/9seconds/mtg/conntypes"
-)
-
-type statsJSON struct {
-	Connections statsJSONConnections `json:"connections"`
-	Traffic     statsJSONTraffic     `json:"traffic"`
-	Uptime      statsJSONUptime      `json:"uptime"`
-	Crashes     uint32               `json:"crashes"`
-	AntiReplays uint32               `json:"anti_replay_detected"`
-}
-
-type statsBaseJSONConnections struct {
-	All          statsJSONConnectionType `json:"all"`
-	Abridged     statsJSONConnectionType `json:"abridged"`
-	Intermediate statsJSONConnectionType `json:"intermediate"`
-	Secured      statsJSONConnectionType `json:"secured"`
-}
-
-type statsJSONConnections struct {
-	statsBaseJSONConnections
-}
-
-type statsJSONConnectionType struct {
-	IPv4 uint32 `json:"ipv4"`
-	IPv6 uint32 `json:"ipv6"`
-}
-
-func (c statsJSONConnections) MarshalJSON() ([]byte, error) {
-	c.All.IPv4 = c.Abridged.IPv4 + c.Intermediate.IPv4 + c.Secured.IPv4
-	c.All.IPv6 = c.Abridged.IPv6 + c.Intermediate.IPv6 + c.Secured.IPv6
-
-	return json.Marshal(c.statsBaseJSONConnections)
-}
-
-type statsJSONTraffic struct {
-	Ingress uint64 `json:"ingress"`
-	Egress  uint64 `json:"egress"`
-}
-
-type statsJSONUptime time.Time
-
-func (s statsJSONUptime) MarshalJSON() ([]byte, error) {
-	seconds := strconv.Itoa(int(time.Since(time.Time(s)).Seconds()))
-	return []byte(seconds), nil
-}
-
-func (s *statsJSON) IngressTraffic(traffic int) {
-	atomic.AddUint64(&s.Traffic.Ingress, uint64(traffic))
-}
-
-func (s *statsJSON) EgressTraffic(traffic int) {
-	atomic.AddUint64(&s.Traffic.Egress, uint64(traffic))
-}
-
-func (s *statsJSON) ClientConnected(connectionType conntypes.ConnectionType, addr *net.TCPAddr) {
-	s.changeConnections(connectionType, addr, 1)
-}
-
-func (s *statsJSON) ClientDisconnected(connectionType conntypes.ConnectionType, addr *net.TCPAddr) {
-	s.changeConnections(connectionType, addr, ^uint32(0))
-}
-
-func (s *statsJSON) changeConnections(connectionType conntypes.ConnectionType, addr *net.TCPAddr, value uint32) {
-	var connections *statsJSONConnectionType
-
-	switch connectionType {
-	case conntypes.ConnectionTypeAbridged:
-		connections = &s.Connections.Abridged
-	case conntypes.ConnectionTypeSecure:
-		connections = &s.Connections.Secured
-	default:
-		connections = &s.Connections.Intermediate
-	}
-
-	if addr.IP.To4() != nil {
-		atomic.AddUint32(&connections.IPv4, value)
-	} else {
-		atomic.AddUint32(&connections.IPv6, value)
-	}
-}
-
-func (s *statsJSON) Crash() {
-	atomic.AddUint32(&s.Crashes, 1)
-}
-
-func (s *statsJSON) AntiReplayDetected() {
-	atomic.AddUint32(&s.AntiReplays, 1)
-}
-
-func newStatsJSON(mux *http.ServeMux) Stats {
-	instance := &statsJSON{
-		Uptime: statsJSONUptime(time.Now()),
-	}
-	logger := zap.S().Named("stats")
-
-	mux.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
-		w.Header().Set("Content-Type", "application/json")
-		first, err := json.Marshal(instance)
-		if err != nil {
-			logger.Errorw("Cannot encode json", "error", err)
-			http.Error(w, "Internal server error", http.StatusServiceUnavailable)
-			return
-		}
-
-		interim := map[string]interface{}{}
-		if err := json.Unmarshal(first, &interim); err != nil {
-			panic(err)
-		}
-
-		encoder := json.NewEncoder(w)
-		encoder.SetEscapeHTML(false)
-		encoder.SetIndent("", "  ")
-		if err := encoder.Encode(interim); err != nil {
-			logger.Errorw("Cannot encode json", "error", err)
-		}
-	})
-
-	return instance
-}

stats/stats_prometheus.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"strconv"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
@@ -13,10 +14,11 @@ import (
 )
 
 type statsPrometheus struct {
-	connections *prometheus.GaugeVec
+	connections         *prometheus.GaugeVec
-	traffic     *prometheus.GaugeVec
+	telegramConnections *prometheus.GaugeVec
-	crashes     prometheus.Gauge
+	traffic             *prometheus.GaugeVec
-	antiReplays prometheus.Gauge
+	crashes             prometheus.Gauge
+	antiReplays         prometheus.Counter
 }
 
 func (s *statsPrometheus) IngressTraffic(traffic int) {
@@ -38,18 +40,39 @@ func (s *statsPrometheus) ClientDisconnected(connectionType conntypes.Connection
 func (s *statsPrometheus) changeConnections(connectionType conntypes.ConnectionType,
 	addr *net.TCPAddr,
 	increment float64) {
-	var labels [2]string
+	labels := [...]string{
+		"intermediate",
+		"ipv4",
+	}
 
 	switch connectionType {
 	case conntypes.ConnectionTypeAbridged:
 		labels[0] = "abridged"
 	case conntypes.ConnectionTypeSecure:
 		labels[0] = "secured"
-	default:
-		labels[0] = "intermediate"
 	}
 
-	labels[1] = "ipv4"
+	if addr.IP.To4() == nil {
+		labels[1] = "ipv6" // nolint: goconst
+	}
+
+	s.connections.WithLabelValues(labels[:]...).Add(increment)
+}
+
+func (s *statsPrometheus) TelegramConnected(dc conntypes.DC, addr *net.TCPAddr) {
+	s.changeTelegramConnections(dc, addr, 1.0)
+}
+
+func (s *statsPrometheus) TelegramDisconnected(dc conntypes.DC, addr *net.TCPAddr) {
+	s.changeTelegramConnections(dc, addr, -1.0)
+}
+
+func (s *statsPrometheus) changeTelegramConnections(dc conntypes.DC, addr *net.TCPAddr, increment float64) {
+	labels := [...]string{
+		strconv.Itoa(int(dc)),
+		"ipv4",
+	}
+
 	if addr.IP.To4() == nil {
 		labels[1] = "ipv6"
 	}
@@ -65,26 +88,32 @@ func (s *statsPrometheus) AntiReplayDetected() {
 	s.antiReplays.Inc()
 }
 
-func newStatsPrometheus(mux *http.ServeMux) (Stats, error) {
+func newStatsPrometheus(mux *http.ServeMux) (Interface, error) {
-	registry := prometheus.NewRegistry()
+	registry := prometheus.NewPedanticRegistry()
+
 	instance := &statsPrometheus{
 		connections: prometheus.NewGaugeVec(prometheus.GaugeOpts{
-			Namespace: config.C.PrometheusStats.Prefix,
+			Namespace: config.C.StatsNamespace,
 			Name:      "connections",
-			Help:      "Current number of connections to the proxy.",
+			Help:      "Current number of client connections to the proxy.",
 		}, []string{"type", "protocol"}),
+		telegramConnections: prometheus.NewGaugeVec(prometheus.GaugeOpts{
+			Namespace: config.C.StatsNamespace,
+			Name:      "telegram_connections",
+			Help:      "Current number of telegram connections established by this proxy.",
+		}, []string{"dc", "protocol"}),
 		traffic: prometheus.NewGaugeVec(prometheus.GaugeOpts{
-			Namespace: config.C.PrometheusStats.Prefix,
+			Namespace: config.C.StatsNamespace,
 			Name:      "traffic",
 			Help:      "Traffic passed through the proxy in bytes.",
 		}, []string{"direction"}),
 		crashes: prometheus.NewGauge(prometheus.GaugeOpts{
-			Namespace: config.C.PrometheusStats.Prefix,
+			Namespace: config.C.StatsNamespace,
 			Name:      "crashes",
 			Help:      "How many crashes happened.",
 		}),
-		antiReplays: prometheus.NewGauge(prometheus.GaugeOpts{
+		antiReplays: prometheus.NewCounter(prometheus.CounterOpts{
-			Namespace: config.C.PrometheusStats.Prefix,
+			Namespace: config.C.StatsNamespace,
 			Name:      "anti_replays",
 			Help:      "How many anti replay attacks were prevented.",
 		}),
@@ -104,7 +133,7 @@ func newStatsPrometheus(mux *http.ServeMux) (Stats, error) {
 	}
 
 	handler := promhttp.HandlerFor(registry, promhttp.HandlerOpts{})
-	mux.Handle("/prometheus", handler)
+	mux.Handle("/", handler)
 
 	return instance, nil
 }

stats/stats_statsd.go

@@ -3,6 +3,7 @@ package stats
 import (
 	"fmt"
 	"net"
+	"strconv"
 	"strings"
 
 	"gopkg.in/alexcesaro/statsd.v2"
@@ -32,19 +33,41 @@ func (s *statsStatsd) ClientDisconnected(connectionType conntypes.ConnectionType
 }
 
 func (s *statsStatsd) changeConnections(connectionType conntypes.ConnectionType, addr *net.TCPAddr, value int) {
-	var labels [3]string
+	labels := [...]string{
+		"connections",
+		"intermediate",
+		"ipv4",
+	}
 
-	labels[0] = "connections"
 	switch connectionType {
 	case conntypes.ConnectionTypeAbridged:
 		labels[1] = "abridged"
 	case conntypes.ConnectionTypeSecure:
 		labels[1] = "secured"
-	default:
-		labels[1] = "intermediate"
 	}
 
-	labels[2] = "ipv4"
+	if addr.IP.To4() == nil {
+		labels[2] = "ipv6"
+	}
+
+	s.client.Count(strings.Join(labels[:], "."), value)
+}
+
+func (s *statsStatsd) TelegramConnected(dc conntypes.DC, addr *net.TCPAddr) {
+	s.changeTelegramConnections(dc, addr, 1)
+}
+
+func (s *statsStatsd) TelegramDisconnected(dc conntypes.DC, addr *net.TCPAddr) {
+	s.changeTelegramConnections(dc, addr, -1)
+}
+
+func (s *statsStatsd) changeTelegramConnections(dc conntypes.DC, addr *net.TCPAddr, value int) {
+	labels := [...]string{
+		"telegram",
+		strconv.Itoa(int(dc)),
+		"ipv4",
+	}
+
 	if addr.IP.To4() == nil {
 		labels[2] = "ipv6"
 	}
@@ -60,17 +83,17 @@ func (s *statsStatsd) AntiReplayDetected() {
 	s.client.Increment("anti_replays")
 }
 
-func newStatsStatsd() (Stats, error) {
+func newStatsStatsd() (Interface, error) {
 	options := []statsd.Option{
-		statsd.Prefix(config.C.StatsdStats.Prefix),
+		statsd.Prefix(config.C.StatsNamespace),
-		statsd.Network(config.C.StatsdStats.Addr.Network()),
+		statsd.Network(config.C.StatsdNetwork),
-		statsd.Address(config.C.StatsdStats.Addr.String()),
+		statsd.Address(config.C.StatsBind.String()),
-		statsd.TagsFormat(config.C.StatsdStats.TagsFormat),
+		statsd.TagsFormat(config.C.StatsdTagsFormat),
 	}
 
-	if len(config.C.StatsdStats.Tags) > 0 {
+	if len(config.C.StatsdTags) > 0 {
-		tags := make([]string, len(config.C.StatsdStats.Tags)*2)
+		tags := make([]string, len(config.C.StatsdTags)*2)
-		for k, v := range config.C.StatsdStats.Tags {
+		for k, v := range config.C.StatsdTags {
 			tags = append(tags, k, v)
 		}
 		options = append(options, statsd.Tags(tags...))

telegram/base.go

@@ -47,7 +47,7 @@ func (b *baseTelegram) dial(dc conntypes.DC,
 		return nil, fmt.Errorf("cannot initialize tcp socket: %w", err)
 	}
 
-	return wrappers.NewTelegramConn(conn), nil
+	return wrappers.NewTelegramConn(dc, conn), nil
 }
 
 func (b *baseTelegram) chooseAddress(addresses map[conntypes.DC][]string,

utils/init_tcp.go

@@ -13,10 +13,10 @@ func InitTCP(conn net.Conn) error {
 	if err := tcpConn.SetNoDelay(true); err != nil {
 		return fmt.Errorf("cannot set TCP_NO_DELAY: %w", err)
 	}
-	if err := tcpConn.SetReadBuffer(config.C.BufferSize.Read); err != nil {
+	if err := tcpConn.SetReadBuffer(config.C.ReadBuffer); err != nil {
 		return fmt.Errorf("cannot set read buffer size: %w", err)
 	}
-	if err := tcpConn.SetWriteBuffer(config.C.BufferSize.Write); err != nil {
+	if err := tcpConn.SetWriteBuffer(config.C.WriteBuffer); err != nil {
 		return fmt.Errorf("cannot set write buffer size: %w", err)
 	}
 

wrappers/common.go

@@ -0,0 +1,21 @@
+package wrappers
+
+import (
+	"net"
+
+	"github.com/9seconds/mtg/conntypes"
+)
+
+func NewClientConn(parent net.Conn, connID conntypes.ConnID) conntypes.StreamReadWriteCloser {
+	conn := newConn(parent, connID, connPurposeClient)
+	conn = NewTrafficStats(conn)
+
+	return conn
+}
+
+func NewTelegramConn(dc conntypes.DC, parent net.Conn) conntypes.StreamReadWriteCloser {
+	conn := newConn(parent, conntypes.ConnID{}, connPurposeTelegram)
+	conn = NewTelegramStats(dc, conn)
+
+	return conn
+}

wrappers/conn.go

@@ -91,11 +91,11 @@ func newConn(parent net.Conn,
 	localAddr := *parent.LocalAddr().(*net.TCPAddr)
 
 	if parent.RemoteAddr().(*net.TCPAddr).IP.To4() != nil {
-		if config.C.PublicIPv4Addr.IP != nil {
+		if config.C.PublicIPv4.IP != nil {
-			localAddr.IP = config.C.PublicIPv4Addr.IP
+			localAddr.IP = config.C.PublicIPv4.IP
 		}
-	} else if config.C.PublicIPv6Addr.IP != nil {
+	} else if config.C.PublicIPv6.IP != nil {
-		localAddr.IP = config.C.PublicIPv6Addr.IP
+		localAddr.IP = config.C.PublicIPv6.IP
 	}
 
 	logger := zap.S().With(
@@ -115,12 +115,3 @@ func newConn(parent net.Conn,
 		localAddr:  &localAddr,
 	}
 }
-
-func NewClientConn(parent net.Conn,
-	connID conntypes.ConnID) conntypes.StreamReadWriteCloser {
-	return newConn(parent, connID, connPurposeClient)
-}
-
-func NewTelegramConn(parent net.Conn) conntypes.StreamReadWriteCloser {
-	return newConn(parent, conntypes.ConnID{}, connPurposeTelegram)
-}

wrappers/stats.go

@@ -1,67 +0,0 @@
-package wrappers
-
-import (
-	"net"
-	"time"
-
-	"go.uber.org/zap"
-
-	"github.com/9seconds/mtg/conntypes"
-	"github.com/9seconds/mtg/stats"
-)
-
-type wrapperStats struct {
-	parent conntypes.StreamReadWriteCloser
-}
-
-func (w *wrapperStats) Write(p []byte) (int, error) {
-	n, err := w.parent.Write(p)
-	stats.S.EgressTraffic(n)
-
-	return n, err
-}
-
-func (w *wrapperStats) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
-	n, err := w.parent.WriteTimeout(p, timeout)
-	stats.S.EgressTraffic(n)
-
-	return n, err
-}
-
-func (w *wrapperStats) Read(p []byte) (int, error) {
-	n, err := w.parent.Read(p)
-	stats.S.IngressTraffic(n)
-
-	return n, err
-}
-
-func (w *wrapperStats) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
-	n, err := w.parent.ReadTimeout(p, timeout)
-	stats.S.IngressTraffic(n)
-
-	return n, err
-}
-
-func (w *wrapperStats) Conn() net.Conn {
-	return w.parent.Conn()
-}
-
-func (w *wrapperStats) Logger() *zap.SugaredLogger {
-	return w.parent.Logger().Named("traffic")
-}
-
-func (w *wrapperStats) LocalAddr() *net.TCPAddr {
-	return w.parent.LocalAddr()
-}
-
-func (w *wrapperStats) RemoteAddr() *net.TCPAddr {
-	return w.parent.RemoteAddr()
-}
-
-func (w *wrapperStats) Close() error {
-	return w.parent.Close()
-}
-
-func NewTraffic(parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser {
-	return &wrapperStats{parent}
-}

wrappers/stats_telegram.go

@@ -0,0 +1,70 @@
+package wrappers
+
+import (
+	"net"
+	"sync"
+	"time"
+
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/stats"
+	"go.uber.org/zap"
+)
+
+type wrapperTelegramStats struct {
+	parent conntypes.StreamReadWriteCloser
+	dc     conntypes.DC
+	once   sync.Once
+}
+
+func (w *wrapperTelegramStats) Write(p []byte) (int, error) {
+	return w.parent.Write(p)
+}
+
+func (w *wrapperTelegramStats) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
+	return w.parent.WriteTimeout(p, timeout)
+}
+
+func (w *wrapperTelegramStats) Read(p []byte) (int, error) {
+	return w.parent.Read(p)
+}
+
+func (w *wrapperTelegramStats) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
+	return w.parent.ReadTimeout(p, timeout)
+}
+
+func (w *wrapperTelegramStats) Conn() net.Conn {
+	return w.parent.Conn()
+}
+
+func (w *wrapperTelegramStats) Logger() *zap.SugaredLogger {
+	return w.parent.Logger().Named("stats-telegram")
+}
+
+func (w *wrapperTelegramStats) LocalAddr() *net.TCPAddr {
+	return w.parent.LocalAddr()
+}
+
+func (w *wrapperTelegramStats) RemoteAddr() *net.TCPAddr {
+	return w.parent.RemoteAddr()
+}
+
+func (w *wrapperTelegramStats) Close() error {
+	var err error
+
+	w.once.Do(func() {
+		err = w.parent.Close()
+		stats.Stats.TelegramDisconnected(w.dc, w.RemoteAddr())
+	})
+
+	return err
+}
+
+func NewTelegramStats(dc conntypes.DC, parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser {
+	conn := &wrapperTelegramStats{
+		parent: parent,
+		dc:     dc,
+	}
+	stats.Stats.TelegramConnected(dc, parent.RemoteAddr())
+
+	return conn
+}

wrappers/stats_traffic.go

@@ -0,0 +1,67 @@
+package wrappers
+
+import (
+	"net"
+	"time"
+
+	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/conntypes"
+	"github.com/9seconds/mtg/stats"
+)
+
+type wrapperTrafficStats struct {
+	parent conntypes.StreamReadWriteCloser
+}
+
+func (w *wrapperTrafficStats) Write(p []byte) (int, error) {
+	n, err := w.parent.Write(p)
+	stats.Stats.EgressTraffic(n)
+
+	return n, err
+}
+
+func (w *wrapperTrafficStats) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
+	n, err := w.parent.WriteTimeout(p, timeout)
+	stats.Stats.EgressTraffic(n)
+
+	return n, err
+}
+
+func (w *wrapperTrafficStats) Read(p []byte) (int, error) {
+	n, err := w.parent.Read(p)
+	stats.Stats.IngressTraffic(n)
+
+	return n, err
+}
+
+func (w *wrapperTrafficStats) ReadTimeout(p []byte, timeout time.Duration) (int, error) {
+	n, err := w.parent.ReadTimeout(p, timeout)
+	stats.Stats.IngressTraffic(n)
+
+	return n, err
+}
+
+func (w *wrapperTrafficStats) Conn() net.Conn {
+	return w.parent.Conn()
+}
+
+func (w *wrapperTrafficStats) Logger() *zap.SugaredLogger {
+	return w.parent.Logger().Named("stats-traffic")
+}
+
+func (w *wrapperTrafficStats) LocalAddr() *net.TCPAddr {
+	return w.parent.LocalAddr()
+}
+
+func (w *wrapperTrafficStats) RemoteAddr() *net.TCPAddr {
+	return w.parent.RemoteAddr()
+}
+
+func (w *wrapperTrafficStats) Close() error {
+	return w.parent.Close()
+}
+
+func NewTrafficStats(parent conntypes.StreamReadWriteCloser) conntypes.StreamReadWriteCloser {
+	return &wrapperTrafficStats{parent}
+}