fix: use goroutine instead of SetReadDeadline to try to fix compatibility

This commit is contained in:
mzz2017 2023-04-12 21:45:50 +08:00
parent 5d62ce3d22
commit 4efed17ff9
2 changed files with 71 additions and 42 deletions

View File

@ -6,26 +6,30 @@
package sniffing
import (
"errors"
"io"
"net"
"sync"
"time"
)
type Sniffer struct {
r io.Reader
// Stream
stream bool
r io.Reader
dataReady chan struct{}
dataError error
// Common
buf []byte
bufAt int
stream bool
readMu sync.Mutex
}
func NewStreamSniffer(r io.Reader, bufSize int) *Sniffer {
s := &Sniffer{
r: r,
buf: make([]byte, bufSize),
stream: true,
r: r,
buf: make([]byte, bufSize),
stream: true,
dataReady: make(chan struct{}),
}
return s
}
@ -40,38 +44,7 @@ func NewPacketSniffer(data []byte) *Sniffer {
type sniff func() (d string, err error)
func (s *Sniffer) SniffTcp() (d string, err error) {
s.readMu.Lock()
defer s.readMu.Unlock()
if s.stream {
r, isConn := s.r.(net.Conn)
if isConn {
// Set timeout.
r.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
}
n, err := r.Read(s.buf)
if isConn {
// Recover.
r.SetReadDeadline(time.Time{})
}
s.buf = s.buf[:n]
if err != nil {
var netError net.Error
if isConn && errors.As(err, &netError) && netError.Timeout() {
goto sniff
}
return "", err
}
}
sniff:
if len(s.buf) == 0 {
return "", NotApplicableError
}
sniffs := []sniff{
// Most sniffable traffic is TLS, thus we sniff it first.
s.SniffTls,
s.SniffHttp,
}
func sniffGroup(sniffs []sniff) (d string, err error) {
for _, sniffer := range sniffs {
d, err = sniffer()
if err == nil {
@ -84,9 +57,65 @@ sniff:
return "", NotApplicableError
}
func (s *Sniffer) Read(p []byte) (n int, err error) {
func (s *Sniffer) SniffTcp() (d string, err error) {
s.readMu.Lock()
defer s.readMu.Unlock()
if s.stream {
go func() {
n, err := s.r.Read(s.buf)
s.buf = s.buf[:n]
if err != nil {
s.dataError = err
}
close(s.dataReady)
}()
// Waiting 100ms for data.
select {
case <-time.After(100 * time.Millisecond):
return "", NotApplicableError
case <-s.dataReady:
if s.dataError != nil {
return "", s.dataError
}
}
} else {
close(s.dataReady)
}
if len(s.buf) == 0 {
return "", NotApplicableError
}
return sniffGroup([]sniff{
// Most sniffable traffic is TLS, thus we sniff it first.
s.SniffTls,
s.SniffHttp,
})
}
func (s *Sniffer) SniffUdp() (d string, err error) {
s.readMu.Lock()
defer s.readMu.Unlock()
if len(s.buf) == 0 {
return "", NotApplicableError
}
return sniffGroup([]sniff{
s.SniffQuic,
})
}
func (s *Sniffer) Read(p []byte) (n int, err error) {
<-s.dataReady
if s.dataError != nil {
return 0, s.dataError
}
s.readMu.Lock()
defer s.readMu.Unlock()
if s.buf != nil && s.bufAt < len(s.buf) {
// Read buf first.
n = copy(p, s.buf[s.bufAt:])

View File

@ -127,9 +127,9 @@ func (c *ControlPlane) handlePkt(lConn *net.UDPConn, data []byte, src, pktDst, r
// We should cache DNS records and set record TTL to 0, in order to monitor the dns req and resp in real time.
isDns := dnsMessage != nil
if !isDns {
// Sniff Quic
// Sniff Quic, ...
sniffer := sniffing.NewPacketSniffer(data)
domain, err = sniffer.SniffQuic()
domain, err = sniffer.SniffUdp()
if err != nil && !sniffing.IsSniffingError(err) {
sniffer.Close()
return err