From 6b68e7433581a9a23e04065762319063df54c2ac Mon Sep 17 00:00:00 2001
From: mzz2017 <2017@duck.com>
Date: Wed, 15 Feb 2023 12:54:25 +0800
Subject: [PATCH] optimize: sniffer

---
 component/sniffing/http.go                    | 19 ++++++++++++++++---
 .../sniffing/internal/quicutils/relocation.go |  3 +++
 component/sniffing/quic.go                    | 10 +++++++++-
 component/sniffing/tls.go                     |  2 +-
 4 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/component/sniffing/http.go b/component/sniffing/http.go
index 3cd37dd..88ee2df 100644
--- a/component/sniffing/http.go
+++ b/component/sniffing/http.go
@@ -5,12 +5,21 @@
 
 package sniffing
 
-import "bytes"
+import (
+	"bytes"
+	"unicode"
+)
 
 func (s *Sniffer) SniffHttp() (d string, err error) {
+	// First byte should be printable.
+	if len(s.buf) == 0 || !unicode.IsPrint(rune(s.buf[0])) {
+		return "", NotApplicableError
+	}
+
+	// Search method.
 	search := s.buf
-	if len(search) > 20 {
-		search = search[:20]
+	if len(search) > 12 {
+		search = search[:12]
 	}
 	method, _, found := bytes.Cut(search, []byte(" "))
 	if !found {
@@ -21,6 +30,10 @@ func (s *Sniffer) SniffHttp() (d string, err error) {
 	default:
 		return "", NotApplicableError
 	}
+
+	// Now we assume it is an HTTP packet. We should not return NotApplicableError after here.
+
+	// Search Host.
 	search = s.buf
 	prefix := []byte("Host: ")
 	_, afterHostKey, found := bytes.Cut(search, prefix)
diff --git a/component/sniffing/internal/quicutils/relocation.go b/component/sniffing/internal/quicutils/relocation.go
index 52128c5..d474d95 100644
--- a/component/sniffing/internal/quicutils/relocation.go
+++ b/component/sniffing/internal/quicutils/relocation.go
@@ -8,6 +8,7 @@ package quicutils
 import (
 	"fmt"
 	"github.com/mzz2017/softwind/pool"
+	"io/fs"
 	"sort"
 )
 
@@ -124,6 +125,8 @@ func ExtractCryptoFrameOffset(remainder []byte, transportOffset int) (offset *Cr
 			UpperAppOffset: int(offset),
 			Data:           remainder[nextField : nextField+int(length)],
 		}, nextField + int(length), nil
+	case Quic_FrameType_ConnectionClose, Quic_FrameType_ConnectionClose2:
+		return nil, 0, fmt.Errorf("connection closed: %w", fs.ErrClosed)
 	default:
 		return nil, 0, fmt.Errorf("%w: %v", UnknownFrameTypeError, frameType)
 	}
diff --git a/component/sniffing/quic.go b/component/sniffing/quic.go
index 25d9247..a96851b 100644
--- a/component/sniffing/quic.go
+++ b/component/sniffing/quic.go
@@ -9,6 +9,7 @@ import (
 	"errors"
 	"github.com/mzz2017/softwind/pool"
 	"github.com/v2rayA/dae/component/sniffing/internal/quicutils"
+	"io/fs"
 )
 
 const (
@@ -54,6 +55,10 @@ func (s *Sniffer) SniffQuic() (d string, err error) {
 			}
 			return "", err
 		}
+		if errors.Is(err, fs.ErrClosed) {
+			// ConnectionClose sniffed.
+			return "", NotFoundError
+		}
 		// Error is not NotApplicableError, should be quic block.
 		isQuic = true
 		if len(nextBlock) == 0 {
@@ -127,10 +132,10 @@ func sniffQuicBlock(buf []byte) (d string, next []byte, err error) {
 	firstByte := header[0]
 	rawPacketNumber := pool.Get(quicutils.MaxPacketNumberLength)
 	copy(rawPacketNumber, header[boundary-quicutils.MaxPacketNumberLength:])
-	defer pool.Put(rawPacketNumber)
 	defer func() {
 		header[0] = firstByte
 		copy(header[boundary-quicutils.MaxPacketNumberLength:], rawPacketNumber)
+		pool.Put(rawPacketNumber)
 	}()
 	plaintext, err := quicutils.DecryptQuicFromPool_(header, blockEnd, destConnId)
 	if err != nil {
@@ -141,6 +146,9 @@ func sniffQuicBlock(buf []byte) (d string, next []byte, err error) {
 	// After here, we should not return NotApplicableError.
 	// And we should return nextFrame.
 	if d, err = extractSniFromQuicPayload(plaintext); err != nil {
+		if errors.Is(err, fs.ErrClosed) {
+			return "", nil, err
+		}
 		return "", buf[blockEnd:], NotFoundError
 	}
 	return d, buf[blockEnd:], nil
diff --git a/component/sniffing/tls.go b/component/sniffing/tls.go
index f2952d9..22c6b20 100644
--- a/component/sniffing/tls.go
+++ b/component/sniffing/tls.go
@@ -103,7 +103,7 @@ func findSniExtension(search quicutils.Locator) (string, error) {
 	var b []byte
 	for {
 		if i+4 >= search.Len() {
-			return "", NotApplicableError
+			return "", NotFoundError
 		}
 		b = search.Range(i, i+4)
 		typ := binary.BigEndian.Uint16(b)