Merge pull request #604 from fatedier/http

fix new http no traffic stats, fix #590
This commit is contained in:
fatedier 2018-01-17 23:28:48 +08:00 committed by GitHub
commit 3e62198612
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 45 additions and 1 deletions

View File

@ -412,6 +412,7 @@ func HandleTcpWorkConnection(localInfo *config.LocalSvrConf, proxyPlugin plugin.
err error err error
) )
remote = workConn remote = workConn
defer remote.Close()
if baseInfo.UseEncryption { if baseInfo.UseEncryption {
remote, err = frpIo.WithEncryption(remote, encKey) remote, err = frpIo.WithEncryption(remote, encKey)
if err != nil { 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)) localConn, err := frpNet.ConnectServer("tcp", fmt.Sprintf("%s:%d", localInfo.LocalIp, localInfo.LocalPort))
if err != nil { if err != nil {
workConn.Error("connect to local service [%s:%d] error: %v", localInfo.LocalIp, localInfo.LocalPort, err) workConn.Error("connect to local service [%s:%d] error: %v", localInfo.LocalIp, localInfo.LocalPort, err)
remote.Close()
return return
} }

View File

@ -287,9 +287,18 @@ func (pxy *HttpProxy) GetRealConn() (workConn frpNet.Conn, err error) {
rwc = frpIo.WithCompression(rwc) rwc = frpIo.WithCompression(rwc)
} }
workConn = frpNet.WrapReadWriteCloserToConn(rwc, tmpConn) workConn = frpNet.WrapReadWriteCloserToConn(rwc, tmpConn)
workConn = frpNet.WrapStatsConn(workConn, pxy.updateStatsAfterClosedConn)
StatsOpenConnection(pxy.GetName())
return return
} }
func (pxy *HttpProxy) updateStatsAfterClosedConn(totalRead, totalWrite int64) {
name := pxy.GetName()
StatsCloseConnection(name)
StatsAddTrafficIn(name, totalWrite)
StatsAddTrafficOut(name, totalRead)
}
func (pxy *HttpProxy) Close() { func (pxy *HttpProxy) Close() {
pxy.BaseProxy.Close() pxy.BaseProxy.Close()
for _, closeFn := range pxy.closeFuncs { for _, closeFn := range pxy.closeFuncs {

View File

@ -174,3 +174,38 @@ func (sc *SharedConn) WriteBuff(buffer []byte) (err error) {
_, err = sc.buf.Write(buffer) _, err = sc.buf.Write(buffer)
return err 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
}