From 1aa8cfe036e7e437187ee54dd7d773e4f62ab9a3 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Mon, 9 Jul 2018 11:42:44 +0300
Subject: [PATCH] Stats callbacks

---
 proxy/proxy.go    |  5 +++++
 stats/channels.go |  2 +-
 wrappers/conn.go  | 11 ++++++++---
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/proxy/proxy.go b/proxy/proxy.go
index e0f3a4e..1f2df20 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -12,6 +12,7 @@ import (
 	"github.com/9seconds/mtg/client"
 	"github.com/9seconds/mtg/config"
 	"github.com/9seconds/mtg/mtproto"
+	"github.com/9seconds/mtg/stats"
 	"github.com/9seconds/mtg/telegram"
 	"github.com/9seconds/mtg/wrappers"
 )
@@ -45,6 +46,7 @@ func (p *Proxy) accept(conn net.Conn) {
 		conn.Close()
 
 		if err := recover(); err != nil {
+			stats.NewCrash()
 			log.Errorw("Crash of accept handler", "error", err)
 		}
 	}()
@@ -58,6 +60,9 @@ func (p *Proxy) accept(conn net.Conn) {
 	}
 	defer client.(io.Closer).Close()
 
+	stats.ClientConnected(opts.ConnectionType, client.RemoteAddr())
+	defer stats.ClientDisconnected(opts.ConnectionType, client.RemoteAddr())
+
 	server, err := p.getTelegramConn(opts, connID)
 	if err != nil {
 		log.Errorw("Cannot initialize server connection", "error", err)
diff --git a/stats/channels.go b/stats/channels.go
index 775bd2e..1ece855 100644
--- a/stats/channels.go
+++ b/stats/channels.go
@@ -45,7 +45,7 @@ func connectionManager() {
 	for event := range ConnectionsChan {
 		instance.mutex.RLock()
 
-		isIPv4 := event.addr.IP.To4() == nil
+		isIPv4 := event.addr.IP.To4() != nil
 		var inc uint32 = 1
 		if !event.connected {
 			inc = ^uint32(0)
diff --git a/wrappers/conn.go b/wrappers/conn.go
index ad853af..d770cb4 100644
--- a/wrappers/conn.go
+++ b/wrappers/conn.go
@@ -5,6 +5,8 @@ import (
 	"time"
 
 	"go.uber.org/zap"
+
+	"github.com/9seconds/mtg/stats"
 )
 
 type ConnPurpose uint8
@@ -31,9 +33,10 @@ const (
 )
 
 type Conn struct {
-	connID     string
-	conn       net.Conn
-	logger     *zap.SugaredLogger
+	connID string
+	conn   net.Conn
+	logger *zap.SugaredLogger
+
 	publicIPv4 net.IP
 	publicIPv6 net.IP
 }
@@ -43,6 +46,7 @@ func (c *Conn) Write(p []byte) (int, error) {
 	n, err := c.conn.Write(p)
 
 	c.logger.Debugw("Write to stream", "bytes", n, "error", err)
+	stats.EgressTraffic(n)
 
 	return n, err
 }
@@ -52,6 +56,7 @@ func (c *Conn) Read(p []byte) (int, error) {
 	n, err := c.conn.Read(p)
 
 	c.logger.Debugw("Read from stream", "bytes", n, "error", err)
+	stats.IngressTraffic(n)
 
 	return n, err
 }