mirrors/dae
1
0
Fork 0
mirror of https://github.com/daeuniverse/dae.git synced 2025-01-27 16:11:05 +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
Show Stats 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"
"strings"
"sync"
"time"
)
type ConnSniffer struct {
net.Conn
*Sniffer
sniffer *Sniffer
mu sync.Mutex
}
@ -22,21 +23,26 @@ type ConnSniffer struct {
func NewConnSniffer(conn net.Conn, snifferBufSize int) *ConnSniffer {
s := &ConnSniffer{
Conn: conn,
Sniffer: NewStreamSniffer(conn, snifferBufSize),
sniffer: NewStreamSniffer(conn, snifferBufSize),
}
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) {
s.mu.Lock()
n, err = s.Sniffer.Read(p)
n, err = s.sniffer.Read(p)
s.mu.Unlock()
return n, err
}
func (s *ConnSniffer) Close() (err error) {
var errs []string
if err = s.Sniffer.Close(); err != nil {
if err = s.sniffer.Close(); err != nil {
errs = append(errs, err.Error())
}
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()
// Sniff target domain.
lConn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
sniffer := sniffing.NewConnSniffer(lConn, TcpSniffBufSize)
domain, err := sniffer.SniffTcp()
if err != nil && !sniffing.IsSniffingError(err) {
return err
}
lConn.SetReadDeadline(time.Time{})
// ConnSniffer should be used later, so we cannot close it now.
defer sniffer.Close()