mirrors/dae
1
0
Fork 0
mirror of https://github.com/daeuniverse/dae.git synced 2025-07-04 23:40:30 +07:00
Code Issues Packages Projects Releases Wiki Activity

refactor(sniffer): refine code

Browse Source
This commit is contained in:
mzz2017
2023-04-10 00:37:48 +08:00
parent 02423fabec
commit 8465229e35
2 changed files with 10 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
component/sniffing/conn_sniffer.go
View File

@ -10,11 +10,12 @@ import (
"net" "net"
"strings" "strings"
"sync" "sync"
"time"
) )
type ConnSniffer struct { type ConnSniffer struct {
net.Conn net.Conn
*Sniffer sniffer *Sniffer
mu sync.Mutex mu sync.Mutex
} }
@ -22,21 +23,26 @@ type ConnSniffer struct {
func NewConnSniffer(conn net.Conn, snifferBufSize int) *ConnSniffer { func NewConnSniffer(conn net.Conn, snifferBufSize int) *ConnSniffer {
s := &ConnSniffer{ s := &ConnSniffer{
Conn: conn, Conn: conn,
Sniffer: NewStreamSniffer(conn, snifferBufSize), sniffer: NewStreamSniffer(conn, snifferBufSize),
} }
return s return s
} }
func (s *ConnSniffer) SniffTcp() (d string, err error) {
s.Conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
defer s.Conn.SetReadDeadline(time.Time{})
return s.sniffer.SniffTcp()
}
func (s *ConnSniffer) Read(p []byte) (n int, err error) { func (s *ConnSniffer) Read(p []byte) (n int, err error) {
s.mu.Lock() s.mu.Lock()
n, err = s.Sniffer.Read(p) n, err = s.sniffer.Read(p)
s.mu.Unlock() s.mu.Unlock()
return n, err return n, err
} }
func (s *ConnSniffer) Close() (err error) { func (s *ConnSniffer) Close() (err error) {
var errs []string var errs []string
if err = s.Sniffer.Close(); err != nil { if err = s.sniffer.Close(); err != nil {
errs = append(errs, err.Error()) errs = append(errs, err.Error())
} }
if err = s.Conn.Close(); err != nil { if err = s.Conn.Close(); err != nil {

2
control/tcp.go
View File

@ -29,13 +29,11 @@ func (c *ControlPlane) handleConn(lConn net.Conn) (err error) {
defer lConn.Close() defer lConn.Close()
// Sniff target domain. // Sniff target domain.
lConn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
sniffer := sniffing.NewConnSniffer(lConn, TcpSniffBufSize) sniffer := sniffing.NewConnSniffer(lConn, TcpSniffBufSize)
domain, err := sniffer.SniffTcp() domain, err := sniffer.SniffTcp()
if err != nil && !sniffing.IsSniffingError(err) { if err != nil && !sniffing.IsSniffingError(err) {
return err return err
} }
lConn.SetReadDeadline(time.Time{})
// ConnSniffer should be used later, so we cannot close it now. // ConnSniffer should be used later, so we cannot close it now.
defer sniffer.Close() defer sniffer.Close()