Jason
5 years ago
6 changed files with 112 additions and 107 deletions
@ -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() |
|||
} |
Loading…
Reference in new issue