Refactor code structure for improved readability and maintainability

config.example.standalone.toml

@@ -9,4 +9,5 @@ AllowUsers = '6fe57e3f-e618-4873-ba96-a76adec22ccd,6fe57e3f-e618-4873-ba96-a76ad
 LogFile = '' # can be empty if you don't want to log to file, so the log will be print to stdout
 DebugLevel = 'debug' # debug, info, warn, error
 IntervalSecond = '7200'
-EnableDataUsageMetering = 'true'
+EnableDataUsageMetering = 'true'
+BufferSize = '8192' # buffer size in bytes for WebSocket and TCP/UDP reads

config.example.toml

@@ -10,5 +10,6 @@ LogFile = 'unchain.log' # 日志文件名,可以为空则不记录日志
 DebugLevel = 'debug' # 日志基本debug, info, warn, error
 IntervalSecond = '7200' #主控服务器推送流量数据的间隔,个人模式不关心
 EnableDataUsageMetering = 'true'
+BufferSize = '8192' # 缓冲区大小,用于WebSocket和TCP/UDP读取
 
 

example.env

@@ -4,4 +4,5 @@ ENV REGISTER_URL=https://unchainapi.bob99.workers.dev/api/node
 ENV SUB_ADDRESSES=a.mojocn.com,b.mojocn.com
 ENV ALLOW_USERS=903bcd04-79e7-429c-bf0c-0456c7de9cdc,903bcd04-79e7-429c-bf0c-0456c7de9cd1
 ENV INTERVAL_SECOND=3600
-ENV ENABLE_DATA_USAGE_METERING=true
+ENV ENABLE_DATA_USAGE_METERING=true
+ENV BUFFER_SIZE=8192

global/config.go

@@ -3,7 +3,6 @@ package global
 import (
 	"bytes"
 	"fmt"
-	"github.com/BurntSushi/toml"
 	"log"
 	"log/slog"
 	"os"
@@ -11,6 +10,8 @@ import (
 	"strconv"
 	"strings"
 	"time"
+
+	"github.com/BurntSushi/toml"
 )
 
 type Config struct {
@@ -26,13 +27,11 @@ type Config struct {
 	BuildTime               string `desc:"build time" def:""`                                                                                //optional build time
 	RunAt                   string `desc:"run at" def:""`                                                                                    //optional run at
 	EnableDataUsageMetering string `desc:"enable data usage metering" def:"true"`                                                            //是否开启用户流量统计,使用true 开启用户流量统计,使用false 关闭用户流量统计
+	BufferSize              string `desc:"buffer size in bytes" def:"8192"`                                                                  //缓冲区大小,用于WebSocket和TCP/UDP读取
 }
 
 func (c Config) EnableUsageMetering() bool {
-	if strings.ToLower(c.EnableDataUsageMetering) != "true" {
-		return false
-	}
-	return true
+	return strings.ToLower(c.EnableDataUsageMetering) == "true"
 }
 
 func (c Config) SubHostWithPort() []string {
@@ -114,6 +113,15 @@ func (c Config) ListenPort() int {
 	return int(iv)
 }
 
+func (c Config) GetBufferSize() int {
+	iv, err := strconv.ParseInt(c.BufferSize, 10, 32)
+	if err != nil {
+		log.Println("failed to parse buffer size:", err)
+		return 8192
+	}
+	return int(iv)
+}
+
 var (
 	gitHash   string
 	buildTime string

server/app.go

@@ -14,6 +14,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/gorilla/websocket"
 	"github.com/unchainese/unchain/global"
 )
 
@@ -22,6 +23,8 @@ type App struct {
 	userUsedTrafficKb sync.Map // string -> int64
 	svr               *http.Server
 	exitSignal        chan os.Signal
+	bufferPool        *sync.Pool
+	upGrader          *websocket.Upgrader
 }
 
 func (app *App) httpSvr() {
@@ -42,11 +45,25 @@ func (app *App) httpSvr() {
 }
 
 func NewApp(c *global.Config, sig chan os.Signal) *App {
+	bufferSize := c.GetBufferSize()
 	app := &App{
 		cfg:               c,
 		userUsedTrafficKb: sync.Map{},
 		exitSignal:        sig,
 		svr:               nil,
+		bufferPool: &sync.Pool{
+			New: func() interface{} {
+				return make([]byte, bufferSize)
+			},
+		},
+		upGrader: &websocket.Upgrader{
+			ReadBufferSize:  bufferSize,
+			WriteBufferSize: bufferSize,
+			CheckOrigin: func(r *http.Request) bool {
+				// Allow all connections by default
+				return true
+			},
+		},
 	}
 	for _, userID := range c.UserIDS() {
 		app.userUsedTrafficKb.Store(userID, int64(0))

server/app_ws_vless.go

@@ -20,7 +20,6 @@ import (
 )
 
 const (
-	buffSize          = 8 << 10
 	contentTypeHeader = "Content-Type"
 	contentTypeJSON   = "application/json"
 	upgradeHeader     = "Upgrade"
@@ -28,15 +27,6 @@ const (
 	secWebSocketProto = "sec-websocket-protocol"
 )
 
-var upGrader = websocket.Upgrader{
-	ReadBufferSize:  buffSize,
-	WriteBufferSize: buffSize,
-	CheckOrigin: func(r *http.Request) bool {
-		// Allow all connections by default
-		return true
-	},
-}
-
 func startDstConnection(vd *schema.ProtoVLESS, timeout time.Duration) (net.Conn, []byte, error) {
 	conn, err := net.DialTimeout(vd.DstProtocol, vd.HostPort(), timeout)
 	if err != nil {
@@ -65,7 +55,7 @@ func (app *App) WsVLESS(w http.ResponseWriter, r *http.Request) {
 		log.Println("Error decoding early data:", err)
 	}
 
-	ws, err := upGrader.Upgrade(w, r, nil)
+	ws, err := app.upGrader.Upgrade(w, r, nil)
 	if err != nil {
 		fmt.Println("Error upgrading to websocket:", err)
 		return
@@ -95,9 +85,9 @@ func (app *App) WsVLESS(w http.ResponseWriter, r *http.Request) {
 	sessionTrafficByteN := int64(len(earlyData))
 
 	if vData.DstProtocol == "udp" {
-		sessionTrafficByteN += vlessUDP(ctx, vData, ws)
+		sessionTrafficByteN += app.vlessUDP(ctx, vData, ws)
 	} else if vData.DstProtocol == "tcp" {
-		sessionTrafficByteN += vlessTCP(ctx, vData, ws)
+		sessionTrafficByteN += app.vlessTCP(ctx, vData, ws)
 	} else {
 		log.Println("Error unsupported protocol:", vData.DstProtocol)
 		return
@@ -105,7 +95,7 @@ func (app *App) WsVLESS(w http.ResponseWriter, r *http.Request) {
 	go app.trafficInc(vData.UUID(), sessionTrafficByteN)
 }
 
-func vlessTCP(ctx context.Context, sv *schema.ProtoVLESS, ws *websocket.Conn) int64 {
+func (app *App) vlessTCP(ctx context.Context, sv *schema.ProtoVLESS, ws *websocket.Conn) int64 {
 	logger := sv.Logger()
 	conn, headerVLESS, err := startDstConnection(sv, time.Millisecond*1000)
 	if err != nil {
@@ -162,7 +152,8 @@ func vlessTCP(ctx context.Context, sv *schema.ProtoVLESS, ws *websocket.Conn) in
 		defer wg.Done()
 		defer cancel() // Cancel context if this goroutine exits
 		hasNotSentHeader := true
-		buf := make([]byte, buffSize)
+		buf := app.bufferPool.Get().([]byte)
+		defer app.bufferPool.Put(buf)
 		for {
 			select {
 			case <-ctx.Done():
@@ -196,7 +187,7 @@ func vlessTCP(ctx context.Context, sv *schema.ProtoVLESS, ws *websocket.Conn) in
 }
 
 // vlessUDP handles UDP traffic over VLESS protocol via WebSocket is tested ok
-func vlessUDP(_ context.Context, sv *schema.ProtoVLESS, ws *websocket.Conn) (trafficMeter int64) {
+func (app *App) vlessUDP(_ context.Context, sv *schema.ProtoVLESS, ws *websocket.Conn) (trafficMeter int64) {
 	logger := sv.Logger()
 	conn, headerVLESS, err := startDstConnection(sv, time.Millisecond*1000)
 	if err != nil {
@@ -211,7 +202,8 @@ func vlessUDP(_ context.Context, sv *schema.ProtoVLESS, ws *websocket.Conn) (tra
 		return
 	}
 
-	buf := make([]byte, buffSize)
+	buf := app.bufferPool.Get().([]byte)
+	defer app.bufferPool.Put(buf)
 	n, err := conn.Read(buf)
 	if err != nil {
 		logger.Error("Error reading from UDP connection:", "err", err)

testkit/config.json

@@ -13,6 +13,7 @@
   ],
   "outbounds": [
     {
+      "tag": "proxy",
       "protocol": "vless",
       "settings": {
         "vnext": [
@@ -34,6 +35,66 @@
           "path": "/wsv/v1?ed=2560"
         }
       }
+    },
+    {
+      "tag": "direct",
+      "protocol": "freedom",
+      "settings": {}
+    },
+    {
+      "tag": "block",
+      "protocol": "blackhole",
+      "settings": {
+        "response": {
+          "type": "http"
+        }
+      }
     }
-  ]
+  ],
+  "routing": {
+    "domainStrategy": "IPOnDemand",
+    "rules": [
+      {
+        "type": "field",
+        "outboundTag": "direct",
+        "domain": [
+          "geosite:cn",
+          "geosite:category-ads-all"
+        ]
+      },
+      {
+        "type": "field",
+        "outboundTag": "direct",
+        "ip": [
+          "geoip:cn",
+          "geoip:private"
+        ]
+      },
+      {
+        "type": "field",
+        "outboundTag": "proxy",
+        "network": "tcp,udp"
+      }
+    ]
+  },
+  "dns": {
+    "servers": [
+      {
+        "address": "https://1.1.1.1/dns-query",
+        "domains": [
+          "geosite:geolocation-!cn"
+        ]
+      },
+      {
+        "address": "223.5.5.5",
+        "domains": [
+          "geosite:cn"
+        ],
+        "expectIPs": [
+          "geoip:cn"
+        ]
+      },
+      "8.8.8.8"
+    ]
+  }
 }