From 0c61057551b567d463b60623ed0ab2d06a18260c Mon Sep 17 00:00:00 2001
From: wwqgtxx <wwqgtxx@gmail.com>
Date: Sun, 23 Apr 2023 20:55:42 +0800
Subject: [PATCH] feat: add `statistic` and `only-tcp` options for smux

---
 adapter/outbound/singmux.go | 24 ++++++++++++++++--------
 docs/config.yaml            |  2 ++
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/adapter/outbound/singmux.go b/adapter/outbound/singmux.go
index 315b53c6..4f3a531c 100644
--- a/adapter/outbound/singmux.go
+++ b/adapter/outbound/singmux.go
@@ -16,9 +16,10 @@ import (
 
 type SingMux struct {
 	C.ProxyAdapter
-	base   ProxyBase
-	client *mux.Client
-	dialer *muxSingDialer
+	base    ProxyBase
+	client  *mux.Client
+	dialer  *muxSingDialer
+	onlyTcp bool
 }
 
 type SingMuxOption struct {
@@ -28,6 +29,8 @@ type SingMuxOption struct {
 	MinStreams     int    `proxy:"min-streams,omitempty"`
 	MaxStreams     int    `proxy:"max-streams,omitempty"`
 	Padding        bool   `proxy:"padding,omitempty"`
+	Statistic      bool   `proxy:"statistic,omitempty"`
+	OnlyTcp        bool   `proxy:"only-tcp,omitempty"`
 }
 
 type ProxyBase interface {
@@ -35,19 +38,20 @@ type ProxyBase interface {
 }
 
 type muxSingDialer struct {
-	dialer dialer.Dialer
-	proxy  C.ProxyAdapter
+	dialer    dialer.Dialer
+	proxy     C.ProxyAdapter
+	statistic bool
 }
 
 var _ N.Dialer = (*muxSingDialer)(nil)
 
 func (d *muxSingDialer) DialContext(ctx context.Context, network string, destination M.Socksaddr) (net.Conn, error) {
-	var cDialer C.Dialer = proxydialer.New(d.proxy, d.dialer, false)
+	var cDialer C.Dialer = proxydialer.New(d.proxy, d.dialer, d.statistic)
 	return cDialer.DialContext(ctx, network, destination.String())
 }
 
 func (d *muxSingDialer) ListenPacket(ctx context.Context, destination M.Socksaddr) (net.PacketConn, error) {
-	var cDialer C.Dialer = proxydialer.New(d.proxy, d.dialer, false)
+	var cDialer C.Dialer = proxydialer.New(d.proxy, d.dialer, d.statistic)
 	return cDialer.ListenPacket(ctx, "udp", "", destination.AddrPort())
 }
 
@@ -62,6 +66,9 @@ func (s *SingMux) DialContext(ctx context.Context, metadata *C.Metadata, opts ..
 }
 
 func (s *SingMux) ListenPacketContext(ctx context.Context, metadata *C.Metadata, opts ...dialer.Option) (_ C.PacketConn, err error) {
+	if s.onlyTcp {
+		return s.ProxyAdapter.ListenPacketContext(ctx, metadata, opts...)
+	}
 	options := s.base.DialOptions(opts...)
 	s.dialer.dialer = dialer.NewDialer(options...)
 	pc, err := s.client.ListenPacket(ctx, M.ParseSocksaddr(metadata.RemoteAddress()))
@@ -75,7 +82,7 @@ func (s *SingMux) ListenPacketContext(ctx context.Context, metadata *C.Metadata,
 }
 
 func NewSingMux(option SingMuxOption, proxy C.ProxyAdapter, base ProxyBase) (C.ProxyAdapter, error) {
-	singDialer := &muxSingDialer{dialer: dialer.NewDialer(), proxy: proxy}
+	singDialer := &muxSingDialer{dialer: dialer.NewDialer(), proxy: proxy, statistic: option.Statistic}
 	client, err := mux.NewClient(mux.Options{
 		Context:        context.TODO(),
 		Dialer:         singDialer,
@@ -93,5 +100,6 @@ func NewSingMux(option SingMuxOption, proxy C.ProxyAdapter, base ProxyBase) (C.P
 		base:         base,
 		client:       client,
 		dialer:       singDialer,
+		onlyTcp:      option.OnlyTcp,
 	}, nil
 }
diff --git a/docs/config.yaml b/docs/config.yaml
index 660976e7..e0fe5577 100644
--- a/docs/config.yaml
+++ b/docs/config.yaml
@@ -309,6 +309,8 @@ proxies: # socks5
       # min-streams: 4 # Minimum multiplexed streams in a connection before opening a new connection. Conflict with max-streams.
       # max-streams: 0 # Maximum multiplexed streams in a connection before opening a new connection. Conflict with max-connections and min-streams.
       # padding: false # Enable padding. Requires sing-box server version 1.3-beta9 or later.
+      # statistic: false # 控制是否将底层连接显示在面板中，方面打断底层连接
+      # only-tcp: false # 如果设置为true, smux的设置将不会对udp生效，udp连接会直接走底层协议
   
   - name: "ss2"
     type: ss