From 3412368d206dc93ed5130c2a7ede334d352ddd4e Mon Sep 17 00:00:00 2001
From: Toby <tobyxdd@gmail.com>
Date: Sat, 29 Jun 2024 16:27:57 -0700
Subject: [PATCH] feat: app sniff options

---
 app/cmd/server.go        | 23 ++++++++++++++++++-----
 app/cmd/server_test.go   |  5 +++++
 app/cmd/server_test.yaml |  5 +++++
 extras/sniff/sniff.go    | 11 ++++++++++-
 4 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/app/cmd/server.go b/app/cmd/server.go
index 1a1ee6d..bca4c80 100644
--- a/app/cmd/server.go
+++ b/app/cmd/server.go
@@ -66,6 +66,7 @@ type serverConfig struct {
 	UDPIdleTimeout        time.Duration               `mapstructure:"udpIdleTimeout"`
 	Auth                  serverConfigAuth            `mapstructure:"auth"`
 	Resolver              serverConfigResolver        `mapstructure:"resolver"`
+	Sniff                 serverConfigSniff           `mapstructure:"sniff"`
 	ACL                   serverConfigACL             `mapstructure:"acl"`
 	Outbounds             []serverConfigOutboundEntry `mapstructure:"outbounds"`
 	TrafficStats          serverConfigTrafficStats    `mapstructure:"trafficStats"`
@@ -181,6 +182,12 @@ type serverConfigResolver struct {
 	HTTPS serverConfigResolverHTTPS `mapstructure:"https"`
 }
 
+type serverConfigSniff struct {
+	Enable        bool          `mapstructure:"enable"`
+	Timeout       time.Duration `mapstructure:"timeout"`
+	RewriteDomain bool          `mapstructure:"rewriteDomain"`
+}
+
 type serverConfigACL struct {
 	File              string        `mapstructure:"file"`
 	Inline            []string      `mapstructure:"inline"`
@@ -543,6 +550,16 @@ func serverConfigOutboundHTTPToOutbound(c serverConfigOutboundHTTP) (outbounds.P
 	return outbounds.NewHTTPOutbound(c.URL, c.Insecure)
 }
 
+func (c *serverConfig) fillRequestHook(hyConfig *server.Config) error {
+	if c.Sniff.Enable {
+		hyConfig.RequestHook = &sniff.Sniffer{
+			Timeout:       c.Sniff.Timeout,
+			RewriteDomain: c.Sniff.RewriteDomain,
+		}
+	}
+	return nil
+}
+
 func (c *serverConfig) fillOutboundConfig(hyConfig *server.Config) error {
 	// Resolver, ACL, actual outbound are all implemented through the Outbound interface.
 	// Depending on the config, we build a chain like this:
@@ -823,6 +840,7 @@ func (c *serverConfig) Config() (*server.Config, error) {
 		c.fillConn,
 		c.fillTLSConfig,
 		c.fillQUICConfig,
+		c.fillRequestHook,
 		c.fillOutboundConfig,
 		c.fillBandwidthConfig,
 		c.fillIgnoreClientBandwidth,
@@ -857,11 +875,6 @@ func runServer(cmd *cobra.Command, args []string) {
 		logger.Fatal("failed to load server config", zap.Error(err))
 	}
 
-	hyConfig.RequestHook = &sniff.Sniffer{
-		Timeout:       4 * time.Second,
-		RewriteDomain: false,
-	}
-
 	s, err := server.NewServer(hyConfig)
 	if err != nil {
 		logger.Fatal("failed to initialize server", zap.Error(err))
diff --git a/app/cmd/server_test.go b/app/cmd/server_test.go
index d81a61a..bd46681 100644
--- a/app/cmd/server_test.go
+++ b/app/cmd/server_test.go
@@ -111,6 +111,11 @@ func TestServerConfig(t *testing.T) {
 				Insecure: true,
 			},
 		},
+		Sniff: serverConfigSniff{
+			Enable:        true,
+			Timeout:       1 * time.Second,
+			RewriteDomain: true,
+		},
 		ACL: serverConfigACL{
 			File: "chnroute.txt",
 			Inline: []string{
diff --git a/app/cmd/server_test.yaml b/app/cmd/server_test.yaml
index 86a2dcf..343b0a9 100644
--- a/app/cmd/server_test.yaml
+++ b/app/cmd/server_test.yaml
@@ -83,6 +83,11 @@ resolver:
     sni: real.stuff.net
     insecure: true
 
+sniff:
+  enable: true
+  timeout: 1s
+  rewriteDomain: true
+
 acl:
   file: chnroute.txt
   inline:
diff --git a/extras/sniff/sniff.go b/extras/sniff/sniff.go
index 79e35f6..68b3fbc 100644
--- a/extras/sniff/sniff.go
+++ b/extras/sniff/sniff.go
@@ -15,6 +15,10 @@ import (
 	quicInternal "github.com/apernet/hysteria/extras/v2/sniff/internal/quic"
 )
 
+const (
+	sniffDefaultTimeout = 4 * time.Second
+)
+
 var _ server.RequestHook = (*Sniffer)(nil)
 
 // Sniffer is a server core RequestHook that performs packet inspection and possibly
@@ -62,7 +66,12 @@ func (h *Sniffer) Check(isUDP bool, reqAddr string) bool {
 }
 
 func (h *Sniffer) TCP(stream quic.Stream, reqAddr *string) ([]byte, error) {
-	err := stream.SetReadDeadline(time.Now().Add(h.Timeout))
+	var err error
+	if h.Timeout == 0 {
+		err = stream.SetReadDeadline(time.Now().Add(sniffDefaultTimeout))
+	} else {
+		err = stream.SetReadDeadline(time.Now().Add(h.Timeout))
+	}
 	if err != nil {
 		return nil, err
 	}