From 2ab3475783d370c5622809b30373e1549cd54779 Mon Sep 17 00:00:00 2001
From: Jason <xjasonlyu@gmail.com>
Date: Wed, 7 Aug 2019 20:48:24 +0800
Subject: [PATCH] hope fix a bug

---
 common/stats/session/session.go |  4 ++--
 proxy/socks/tcp.go              |  9 ++-------
 proxy/socks/udp.go              | 12 +++++-------
 3 files changed, 9 insertions(+), 16 deletions(-)

diff --git a/common/stats/session/session.go b/common/stats/session/session.go
index 3871306..ea71728 100644
--- a/common/stats/session/session.go
+++ b/common/stats/session/session.go
@@ -128,9 +128,9 @@ func (s *simpleSessionStater) RemoveSession(key interface{}) {
 			s.completedSessions = s.completedSessions[1:]
 		}
 		s.mux.Unlock()
-
-		s.sessions.Delete(key)
 	} else {
 		log.Warnf("cannot found session key: [%v] ", key)
 	}
+	// delete
+	s.sessions.Delete(key)
 }
diff --git a/proxy/socks/tcp.go b/proxy/socks/tcp.go
index b2552ae..9bb7693 100644
--- a/proxy/socks/tcp.go
+++ b/proxy/socks/tcp.go
@@ -1,7 +1,6 @@
 package socks
 
 import (
-	"fmt"
 	"io"
 	"net"
 	"strconv"
@@ -20,8 +19,6 @@ import (
 type tcpHandler struct {
 	sync.Mutex
 
-	sessKey string
-
 	proxyHost string
 	proxyPort uint16
 
@@ -112,7 +109,7 @@ func (h *tcpHandler) relay(localConn, remoteConn net.Conn, sess *stats.Session)
 	<-upCh // Wait for uplink done.
 
 	if h.sessionStater != nil {
-		h.sessionStater.RemoveSession(h.sessKey)
+		h.sessionStater.RemoveSession(localConn)
 	}
 }
 
@@ -157,8 +154,7 @@ func (h *tcpHandler) Handle(localConn net.Conn, target *net.TCPAddr) error {
 			DownloadBytes: 0,
 			SessionStart:  time.Now(),
 		}
-		h.sessKey = fmt.Sprintf("%s:%s", localConn.LocalAddr().Network(), localConn.LocalAddr().String())
-		h.sessionStater.AddSession(h.sessKey, sess)
+		h.sessionStater.AddSession(localConn, sess)
 	}
 
 	// set keepalive
@@ -169,7 +165,6 @@ func (h *tcpHandler) Handle(localConn net.Conn, target *net.TCPAddr) error {
 	go h.relay(localConn, remoteConn, sess)
 
 	log.Access(process, "proxy", target.Network(), localConn.LocalAddr().String(), targetAddr)
-
 	return nil
 }
 
diff --git a/proxy/socks/udp.go b/proxy/socks/udp.go
index d4b59a8..840c576 100644
--- a/proxy/socks/udp.go
+++ b/proxy/socks/udp.go
@@ -20,11 +20,10 @@ type udpHandler struct {
 	sync.Mutex
 
 	closed  bool
-	sessKey string
+	timeout time.Duration
 
 	proxyHost string
 	proxyPort uint16
-	timeout   time.Duration
 
 	udpConns    map[core.UDPConn]net.PacketConn
 	tcpConns    map[core.UDPConn]net.Conn
@@ -88,7 +87,7 @@ func (h *udpHandler) fetchUDPInput(conn core.UDPConn, input net.PacketConn) {
 		}
 		n, err = conn.WriteFrom(buf[int(3+len(addr)):n], resolvedAddr)
 		if n > 0 && h.sessionStater != nil {
-			if sess := h.sessionStater.GetSession(h.sessKey); sess != nil {
+			if sess := h.sessionStater.GetSession(conn); sess != nil {
 				sess.AddDownloadBytes(int64(n))
 			}
 		}
@@ -196,8 +195,7 @@ func (h *udpHandler) connectInternal(conn core.UDPConn, targetAddr string) error
 				DownloadBytes: 0,
 				SessionStart:  time.Now(),
 			}
-			h.sessKey = fmt.Sprintf("%s:%s", conn.LocalAddr().Network(), conn.LocalAddr().String())
-			h.sessionStater.AddSession(h.sessKey, sess)
+			h.sessionStater.AddSession(conn, sess)
 		}
 		log.Access(process, "proxy", "udp", conn.LocalAddr().String(), targetAddr)
 	}
@@ -224,7 +222,7 @@ func (h *udpHandler) ReceiveTo(conn core.UDPConn, data []byte, addr *net.UDPAddr
 		buf = append(buf, data[:]...)
 		n, err := remoteUDPConn.WriteTo(buf, remoteAddr)
 		if n > 0 && h.sessionStater != nil {
-			if sess := h.sessionStater.GetSession(h.sessKey); sess != nil {
+			if sess := h.sessionStater.GetSession(conn); sess != nil {
 				sess.AddUploadBytes(int64(n))
 			}
 		}
@@ -260,7 +258,7 @@ func (h *udpHandler) Close(conn core.UDPConn) {
 	delete(h.remoteAddrs, conn)
 
 	if h.sessionStater != nil {
-		h.sessionStater.RemoveSession(h.sessKey)
+		h.sessionStater.RemoveSession(conn)
 	}
 
 	h.closed = true