mv session to constant

This commit is contained in:
Jason
2019-08-15 16:19:07 +08:00
parent 6f15e67c5f
commit 1cff4803bd
6 changed files with 112 additions and 107 deletions

View File

@@ -1,6 +1,5 @@
package main package main
import "C"
import ( import (
"flag" "flag"
"fmt" "fmt"

View File

@@ -10,7 +10,6 @@ import (
"sync/atomic" "sync/atomic"
"github.com/xjasonlyu/tun2socks/common/queue" "github.com/xjasonlyu/tun2socks/common/queue"
"github.com/xjasonlyu/tun2socks/component/stats"
C "github.com/xjasonlyu/tun2socks/constant" C "github.com/xjasonlyu/tun2socks/constant"
"github.com/xjasonlyu/tun2socks/log" "github.com/xjasonlyu/tun2socks/log"
) )
@@ -28,7 +27,7 @@ type simpleSessionStater struct {
completedSessionQueue *queue.Queue completedSessionQueue *queue.Queue
} }
func NewSimpleSessionStater() stats.SessionStater { func NewSimpleSessionStater() *simpleSessionStater {
return &simpleSessionStater{ return &simpleSessionStater{
completedSessionQueue: queue.New(maxCompletedSessions), completedSessionQueue: queue.New(maxCompletedSessions),
} }
@@ -36,21 +35,21 @@ func NewSimpleSessionStater() stats.SessionStater {
func (s *simpleSessionStater) sessionStatsHandler(resp http.ResponseWriter, req *http.Request) { func (s *simpleSessionStater) sessionStatsHandler(resp http.ResponseWriter, req *http.Request) {
// Slice of active sessions // Slice of active sessions
var activeSessions []*stats.Session var activeSessions []*C.Session
s.activeSessionMap.Range(func(key, value interface{}) bool { s.activeSessionMap.Range(func(key, value interface{}) bool {
activeSessions = append(activeSessions, value.(*stats.Session)) activeSessions = append(activeSessions, value.(*C.Session))
return true return true
}) })
// Slice of completed sessions // Slice of completed sessions
var completedSessions []*stats.Session var completedSessions []*C.Session
for _, item := range s.completedSessionQueue.Copy() { for _, item := range s.completedSessionQueue.Copy() {
if sess, ok := item.(*stats.Session); ok { if sess, ok := item.(*C.Session); ok {
completedSessions = append(completedSessions, sess) completedSessions = append(completedSessions, sess)
} }
} }
tablePrint := func(w io.Writer, sessions []*stats.Session) { tablePrint := func(w io.Writer, sessions []*C.Session) {
// Sort by session start time. // Sort by session start time.
sort.Slice(sessions, func(i, j int) bool { sort.Slice(sessions, func(i, j int) bool {
return sessions[i].SessionStart.Sub(sessions[j].SessionStart) < 0 return sessions[i].SessionStart.Sub(sessions[j].SessionStart) < 0
@@ -114,13 +113,13 @@ func (s *simpleSessionStater) Stop() error {
return s.server.Close() return s.server.Close()
} }
func (s *simpleSessionStater) AddSession(key interface{}, session *stats.Session) { func (s *simpleSessionStater) AddSession(key interface{}, session *C.Session) {
s.activeSessionMap.Store(key, session) s.activeSessionMap.Store(key, session)
} }
func (s *simpleSessionStater) GetSession(key interface{}) *stats.Session { func (s *simpleSessionStater) GetSession(key interface{}) *C.Session {
if sess, ok := s.activeSessionMap.Load(key); ok { if sess, ok := s.activeSessionMap.Load(key); ok {
return sess.(*stats.Session) return sess.(*C.Session)
} }
return nil return nil
} }

View File

@@ -1,102 +1,13 @@
package stats package stats
import ( import (
"net" C "github.com/xjasonlyu/tun2socks/constant"
"sync"
"sync/atomic"
"time"
) )
type SessionStater interface { type SessionStater interface {
Start() error Start() error
Stop() error Stop() error
AddSession(key interface{}, session *Session) AddSession(key interface{}, session *C.Session)
GetSession(key interface{}) *Session GetSession(key interface{}) *C.Session
RemoveSession(key interface{}) RemoveSession(key interface{})
} }
type Session struct {
Process string
Network string
DialerAddr string
ClientAddr string
TargetAddr string
UploadBytes int64
DownloadBytes int64
SessionStart time.Time
SessionClose time.Time
}
// Track SessionConn
type SessionConn struct {
net.Conn
once sync.Once
session *Session
}
func NewSessionConn(conn net.Conn, session *Session) net.Conn {
return &SessionConn{
Conn: conn,
session: session,
}
}
func (c *SessionConn) Read(b []byte) (n int, err error) {
n, err = c.Conn.Read(b)
if n > 0 {
atomic.AddInt64(&c.session.DownloadBytes, int64(n))
}
return
}
func (c *SessionConn) Write(b []byte) (n int, err error) {
n, err = c.Conn.Write(b)
if n > 0 {
atomic.AddInt64(&c.session.UploadBytes, int64(n))
}
return
}
func (c *SessionConn) Close() error {
c.once.Do(func() {
c.session.SessionClose = time.Now()
})
return c.Conn.Close()
}
// Track SessionPacketConn
type SessionPacketConn struct {
net.PacketConn
once sync.Once
session *Session
}
func NewSessionPacketConn(conn net.PacketConn, session *Session) net.PacketConn {
return &SessionPacketConn{
PacketConn: conn,
session: session,
}
}
func (c *SessionPacketConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
n, addr, err = c.PacketConn.ReadFrom(b)
if n > 0 {
atomic.AddInt64(&c.session.DownloadBytes, int64(n))
}
return
}
func (c *SessionPacketConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
n, err = c.PacketConn.WriteTo(b, addr)
if n > 0 {
atomic.AddInt64(&c.session.UploadBytes, int64(n))
}
return
}
func (c *SessionPacketConn) Close() error {
c.once.Do(func() {
c.session.SessionClose = time.Now()
})
return c.PacketConn.Close()
}

94
constant/session.go Normal file
View File

@@ -0,0 +1,94 @@
package constant
import (
"net"
"sync"
"sync/atomic"
"time"
)
type Session struct {
Process string
Network string
DialerAddr string
ClientAddr string
TargetAddr string
UploadBytes int64
DownloadBytes int64
SessionStart time.Time
SessionClose time.Time
}
// Track SessionConn
type SessionConn struct {
net.Conn
once sync.Once
session *Session
}
func NewSessionConn(conn net.Conn, session *Session) net.Conn {
return &SessionConn{
Conn: conn,
session: session,
}
}
func (c *SessionConn) Read(b []byte) (n int, err error) {
n, err = c.Conn.Read(b)
if n > 0 {
atomic.AddInt64(&c.session.DownloadBytes, int64(n))
}
return
}
func (c *SessionConn) Write(b []byte) (n int, err error) {
n, err = c.Conn.Write(b)
if n > 0 {
atomic.AddInt64(&c.session.UploadBytes, int64(n))
}
return
}
func (c *SessionConn) Close() error {
c.once.Do(func() {
c.session.SessionClose = time.Now()
})
return c.Conn.Close()
}
// Track SessionPacketConn
type SessionPacketConn struct {
net.PacketConn
once sync.Once
session *Session
}
func NewSessionPacketConn(conn net.PacketConn, session *Session) net.PacketConn {
return &SessionPacketConn{
PacketConn: conn,
session: session,
}
}
func (c *SessionPacketConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
n, addr, err = c.PacketConn.ReadFrom(b)
if n > 0 {
atomic.AddInt64(&c.session.DownloadBytes, int64(n))
}
return
}
func (c *SessionPacketConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
n, err = c.PacketConn.WriteTo(b, addr)
if n > 0 {
atomic.AddInt64(&c.session.UploadBytes, int64(n))
}
return
}
func (c *SessionPacketConn) Close() error {
c.once.Do(func() {
c.session.SessionClose = time.Now()
})
return c.PacketConn.Close()
}

View File

@@ -11,6 +11,7 @@ import (
"github.com/xjasonlyu/tun2socks/common/pool" "github.com/xjasonlyu/tun2socks/common/pool"
"github.com/xjasonlyu/tun2socks/component/dns" "github.com/xjasonlyu/tun2socks/component/dns"
"github.com/xjasonlyu/tun2socks/component/stats" "github.com/xjasonlyu/tun2socks/component/stats"
C "github.com/xjasonlyu/tun2socks/constant"
"github.com/xjasonlyu/tun2socks/core" "github.com/xjasonlyu/tun2socks/core"
"github.com/xjasonlyu/tun2socks/log" "github.com/xjasonlyu/tun2socks/log"
) )
@@ -104,7 +105,7 @@ func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
// Get name of the process. // Get name of the process.
var process = lsof.GetProcessName(localConn.LocalAddr()) var process = lsof.GetProcessName(localConn.LocalAddr())
if h.sessionStater != nil { if h.sessionStater != nil {
sess := &stats.Session{ sess := &C.Session{
Process: process, Process: process,
Network: localConn.LocalAddr().Network(), Network: localConn.LocalAddr().Network(),
DialerAddr: remoteConn.LocalAddr().String(), DialerAddr: remoteConn.LocalAddr().String(),
@@ -116,7 +117,7 @@ func (h *tcpHandler) Handle(conn net.Conn, target *net.TCPAddr) error {
} }
h.sessionStater.AddSession(localConn, sess) h.sessionStater.AddSession(localConn, sess)
remoteConn = stats.NewSessionConn(remoteConn, sess) remoteConn = C.NewSessionConn(remoteConn, sess)
} }
// Set keepalive // Set keepalive

View File

@@ -11,6 +11,7 @@ import (
"github.com/xjasonlyu/tun2socks/common/pool" "github.com/xjasonlyu/tun2socks/common/pool"
"github.com/xjasonlyu/tun2socks/component/dns" "github.com/xjasonlyu/tun2socks/component/dns"
"github.com/xjasonlyu/tun2socks/component/stats" "github.com/xjasonlyu/tun2socks/component/stats"
C "github.com/xjasonlyu/tun2socks/constant"
"github.com/xjasonlyu/tun2socks/core" "github.com/xjasonlyu/tun2socks/core"
"github.com/xjasonlyu/tun2socks/log" "github.com/xjasonlyu/tun2socks/log"
) )
@@ -82,7 +83,7 @@ func (h *udpHandler) Connect(conn core.UDPConn, target *net.UDPAddr) error {
// Get name of the process. // Get name of the process.
var process = lsof.GetProcessName(conn.LocalAddr()) var process = lsof.GetProcessName(conn.LocalAddr())
if h.sessionStater != nil { if h.sessionStater != nil {
sess := &stats.Session{ sess := &C.Session{
Process: process, Process: process,
Network: conn.LocalAddr().Network(), Network: conn.LocalAddr().Network(),
DialerAddr: remoteConn.LocalAddr().String(), DialerAddr: remoteConn.LocalAddr().String(),
@@ -94,7 +95,7 @@ func (h *udpHandler) Connect(conn core.UDPConn, target *net.UDPAddr) error {
} }
h.sessionStater.AddSession(conn, sess) h.sessionStater.AddSession(conn, sess)
remoteConn = stats.NewSessionPacketConn(remoteConn, sess) remoteConn = C.NewSessionPacketConn(remoteConn, sess)
} }
h.remoteAddrMap.Store(conn, remoteAddr) h.remoteAddrMap.Store(conn, remoteAddr)