From 4f7dfcdb318abae741677a2ddfde3cf2cba1f66d Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Wed, 17 Jan 2018 23:17:15 +0800
Subject: [PATCH] fix new http no traffic stats, fix #590

---
 client/proxy.go   |  2 +-
 server/proxy.go   |  9 +++++++++
 utils/net/conn.go | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/client/proxy.go b/client/proxy.go
index 005be3a..1802f78 100644
--- a/client/proxy.go
+++ b/client/proxy.go
@@ -412,6 +412,7 @@ func HandleTcpWorkConnection(localInfo *config.LocalSvrConf, proxyPlugin plugin.
 		err    error
 	)
 	remote = workConn
+	defer remote.Close()
 	if baseInfo.UseEncryption {
 		remote, err = frpIo.WithEncryption(remote, encKey)
 		if err != nil {
@@ -433,7 +434,6 @@ func HandleTcpWorkConnection(localInfo *config.LocalSvrConf, proxyPlugin plugin.
 		localConn, err := frpNet.ConnectServer("tcp", fmt.Sprintf("%s:%d", localInfo.LocalIp, localInfo.LocalPort))
 		if err != nil {
 			workConn.Error("connect to local service [%s:%d] error: %v", localInfo.LocalIp, localInfo.LocalPort, err)
-			remote.Close()
 			return
 		}
 
diff --git a/server/proxy.go b/server/proxy.go
index bfb9793..554e818 100644
--- a/server/proxy.go
+++ b/server/proxy.go
@@ -287,9 +287,18 @@ func (pxy *HttpProxy) GetRealConn() (workConn frpNet.Conn, err error) {
 		rwc = frpIo.WithCompression(rwc)
 	}
 	workConn = frpNet.WrapReadWriteCloserToConn(rwc, tmpConn)
+	workConn = frpNet.WrapStatsConn(workConn, pxy.updateStatsAfterClosedConn)
+	StatsOpenConnection(pxy.GetName())
 	return
 }
 
+func (pxy *HttpProxy) updateStatsAfterClosedConn(totalRead, totalWrite int64) {
+	name := pxy.GetName()
+	StatsCloseConnection(name)
+	StatsAddTrafficIn(name, totalWrite)
+	StatsAddTrafficOut(name, totalRead)
+}
+
 func (pxy *HttpProxy) Close() {
 	pxy.BaseProxy.Close()
 	for _, closeFn := range pxy.closeFuncs {
diff --git a/utils/net/conn.go b/utils/net/conn.go
index c1f6f46..78319cc 100644
--- a/utils/net/conn.go
+++ b/utils/net/conn.go
@@ -174,3 +174,38 @@ func (sc *SharedConn) WriteBuff(buffer []byte) (err error) {
 	_, err = sc.buf.Write(buffer)
 	return err
 }
+
+type StatsConn struct {
+	Conn
+
+	totalRead  int64
+	totalWrite int64
+	statsFunc  func(totalRead, totalWrite int64)
+}
+
+func WrapStatsConn(conn Conn, statsFunc func(total, totalWrite int64)) *StatsConn {
+	return &StatsConn{
+		Conn:      conn,
+		statsFunc: statsFunc,
+	}
+}
+
+func (statsConn *StatsConn) Read(p []byte) (n int, err error) {
+	n, err = statsConn.Conn.Read(p)
+	statsConn.totalRead += int64(n)
+	return
+}
+
+func (statsConn *StatsConn) Write(p []byte) (n int, err error) {
+	n, err = statsConn.Conn.Write(p)
+	statsConn.totalWrite += int64(n)
+	return
+}
+
+func (statsConn *StatsConn) Close() (err error) {
+	err = statsConn.Conn.Close()
+	if statsConn.statsFunc != nil {
+		statsConn.statsFunc(statsConn.totalRead, statsConn.totalWrite)
+	}
+	return
+}