From 8465229e35ea66abc7743423758ffbdde3df51b0 Mon Sep 17 00:00:00 2001
From: mzz2017 <2017@duck.com>
Date: Mon, 10 Apr 2023 00:37:48 +0800
Subject: [PATCH] refactor(sniffer): refine code

---
 component/sniffing/conn_sniffer.go | 14 ++++++++++----
 control/tcp.go                     |  2 --
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/component/sniffing/conn_sniffer.go b/component/sniffing/conn_sniffer.go
index 9aae0ff..7df3a13 100644
--- a/component/sniffing/conn_sniffer.go
+++ b/component/sniffing/conn_sniffer.go
@@ -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 {
diff --git a/control/tcp.go b/control/tcp.go
index 41bbe03..4a10b0e 100644
--- a/control/tcp.go
+++ b/control/tcp.go
@@ -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()