udp stats

cmd/tun2socks/main_d.go

@@ -31,7 +31,7 @@ func init() {
 		proxyPort := uint16(proxyAddr.Port)
 
 		proxyTCPHandler := socks.NewTCPHandler(proxyHost, proxyPort, fakeDns, sessionStater)
-		proxyUDPHandler := socks.NewUDPHandler(proxyHost, proxyPort, *args.UdpTimeout, dnsCache, fakeDns)
+		proxyUDPHandler := socks.NewUDPHandler(proxyHost, proxyPort, *args.UdpTimeout, dnsCache, fakeDns, sessionStater)
 
 		sendThrough, err := net.ResolveTCPAddr("tcp", *args.ExceptionSendThrough)
 		if err != nil {

cmd/tun2socks/main_socks.go

@@ -25,6 +25,6 @@ func init() {
 		proxyPort := uint16(proxyAddr.Port)
 
 		core.RegisterTCPConnHandler(socks.NewTCPHandler(proxyHost, proxyPort, fakeDns, sessionStater))
-		core.RegisterUDPConnHandler(socks.NewUDPHandler(proxyHost, proxyPort, *args.UdpTimeout, dnsCache, fakeDns))
+		core.RegisterUDPConnHandler(socks.NewUDPHandler(proxyHost, proxyPort, *args.UdpTimeout, dnsCache, fakeDns, sessionStater))
 	})
 }

common/stats/session/session.go

@@ -33,6 +33,13 @@ func (s *simpleSessionStater) AddSession(key interface{}, session *stats.Session
 	s.sessions.Store(key, session)
 }
 
+func (s *simpleSessionStater) GetSession(key interface{}) *stats.Session {
+	if sess, ok := s.sessions.Load(key); ok {
+		return sess.(*stats.Session)
+	}
+	return nil
+}
+
 func (s *simpleSessionStater) RemoveSession(key interface{}) {
 	if sess, ok := s.sessions.Load(key); ok {
 		s.completedSessions = append(s.completedSessions, *(sess.(*stats.Session)))

common/stats/stats.go

@@ -7,6 +7,7 @@ import (
 
 type SessionStater interface {
 	AddSession(key interface{}, session *Session)
+	GetSession(key interface{}) *Session
 	RemoveSession(key interface{})
 }
 

proxy/socks/udp.go

@@ -11,6 +11,8 @@ import (
 
 	"github.com/eycorsican/go-tun2socks/common/dns"
 	"github.com/eycorsican/go-tun2socks/common/log"
+	"github.com/eycorsican/go-tun2socks/common/lsof"
+	"github.com/eycorsican/go-tun2socks/common/stats"
 	"github.com/eycorsican/go-tun2socks/core"
 )
 
@@ -22,21 +24,24 @@ type udpHandler struct {
 	udpConns    map[core.UDPConn]net.PacketConn
 	tcpConns    map[core.UDPConn]net.Conn
 	remoteAddrs map[core.UDPConn]*net.UDPAddr // UDP relay server addresses
-	dnsCache    dns.DnsCache
-	fakeDns     dns.FakeDns
 	timeout     time.Duration
+
+	dnsCache      dns.DnsCache
+	fakeDns       dns.FakeDns
+	sessionStater stats.SessionStater
 }
 
-func NewUDPHandler(proxyHost string, proxyPort uint16, timeout time.Duration, dnsCache dns.DnsCache, fakeDns dns.FakeDns) core.UDPConnHandler {
+func NewUDPHandler(proxyHost string, proxyPort uint16, timeout time.Duration, dnsCache dns.DnsCache, fakeDns dns.FakeDns, sessionStater stats.SessionStater) core.UDPConnHandler {
 	return &udpHandler{
-		proxyHost:   proxyHost,
-		proxyPort:   proxyPort,
-		udpConns:    make(map[core.UDPConn]net.PacketConn, 8),
-		tcpConns:    make(map[core.UDPConn]net.Conn, 8),
-		remoteAddrs: make(map[core.UDPConn]*net.UDPAddr, 8),
-		dnsCache:    dnsCache,
-		fakeDns:     fakeDns,
-		timeout:     timeout,
+		proxyHost:     proxyHost,
+		proxyPort:     proxyPort,
+		udpConns:      make(map[core.UDPConn]net.PacketConn, 8),
+		tcpConns:      make(map[core.UDPConn]net.Conn, 8),
+		remoteAddrs:   make(map[core.UDPConn]*net.UDPAddr, 8),
+		dnsCache:      dnsCache,
+		fakeDns:       fakeDns,
+		timeout:       timeout,
+		sessionStater: sessionStater,
 	}
 }
 
@@ -79,7 +84,12 @@ func (h *udpHandler) fetchUDPInput(conn core.UDPConn, input net.PacketConn) {
 		if err != nil {
 			return
 		}
-		_, err = conn.WriteFrom(buf[int(3+len(addr)):n], resolvedAddr)
+		n, err = conn.WriteFrom(buf[int(3+len(addr)):n], resolvedAddr)
+		if n > 0 && h.sessionStater != nil {
+			if sess := h.sessionStater.GetSession(conn); sess != nil {
+				sess.AddDownloadBytes(int64(n))
+			}
+		}
 		if err != nil {
 			log.Warnf("write local failed: %v", err)
 			return
@@ -174,9 +184,32 @@ func (h *udpHandler) connectInternal(conn core.UDPConn, dest string) error {
 	h.udpConns[conn] = pc
 	h.remoteAddrs[conn] = resolvedRemoteAddr
 	h.Unlock()
+
 	go h.fetchUDPInput(conn, pc)
+
 	if len(dest) != 0 {
-		log.Access("", "proxy", "udp", conn.LocalAddr().String(), dest)
+		var process string
+		if h.sessionStater != nil {
+			// Get name of the process.
+			localHost, localPortStr, _ := net.SplitHostPort(conn.LocalAddr().String())
+			localPortInt, _ := strconv.Atoi(localPortStr)
+			process, err = lsof.GetCommandNameBySocket(conn.LocalAddr().Network(), localHost, uint16(localPortInt))
+			if err != nil {
+				process = "unknown process"
+			}
+
+			sess := &stats.Session{
+				process,
+				conn.LocalAddr().Network(),
+				conn.LocalAddr().String(),
+				dest,
+				0,
+				0,
+				time.Now(),
+			}
+			h.sessionStater.AddSession(conn, sess)
+		}
+		log.Access(process, "proxy", "udp", conn.LocalAddr().String(), dest)
 	}
 	return nil
 }
@@ -232,7 +265,12 @@ func (h *udpHandler) ReceiveTo(conn core.UDPConn, data []byte, addr *net.UDPAddr
 
 		buf := append([]byte{0, 0, 0}, ParseAddr(dest)...)
 		buf = append(buf, data[:]...)
-		_, err := pc.WriteTo(buf, remoteAddr)
+		n, err := pc.WriteTo(buf, remoteAddr)
+		if n > 0 && h.sessionStater != nil {
+			if sess := h.sessionStater.GetSession(conn); sess != nil {
+				sess.AddUploadBytes(int64(n))
+			}
+		}
 		if err != nil {
 			h.Close(conn)
 			return errors.New(fmt.Sprintf("write remote failed: %v", err))
@@ -259,4 +297,8 @@ func (h *udpHandler) Close(conn core.UDPConn) {
 		delete(h.udpConns, conn)
 	}
 	delete(h.remoteAddrs, conn)
+
+	if h.sessionStater != nil {
+		h.sessionStater.RemoveSession(conn)
+	}
 }