refactor: refactor server configuration and logging structure

src/internal/config/config.go

@@ -1,10 +1,20 @@
 package config
 
-import "flag"
+import (
+	"flag"
+	"fmt"
+)
+
+const (
+	ServerModeSocks5 = "socks5"
+	ServerModeTProxy = "tproxy"
+)
 
 type Config struct {
+	ServerMode           string
 	BindAddr             string
 	Port                 int
+	ListenAddr           string
 	LogLevel             string
 	PayloadUA            string
 	UAPattern            string
@@ -13,15 +23,17 @@ type Config struct {
 
 func Parse() (*Config, bool) {
 	var (
-		bindAddr  string
-		port      int
-		loglevel  string
-		payloadUA string
-		uaPattern string
-		partial   bool
-		showVer   bool
+		serverMode string
+		bindAddr   string
+		port       int
+		loglevel   string
+		payloadUA  string
+		uaPattern  string
+		partial    bool
+		showVer    bool
 	)
 
+	flag.StringVar(&serverMode, "m", ServerModeSocks5, "server mode: socks5 or tproxy (default: socks5)")
 	flag.StringVar(&bindAddr, "b", "127.0.0.1", "bind address (default: 127.0.0.1)")
 	flag.IntVar(&port, "p", 1080, "port")
 	flag.StringVar(&payloadUA, "f", "FFF", "User-Agent")
@@ -32,8 +44,10 @@ func Parse() (*Config, bool) {
 	flag.Parse()
 
 	cfg := &Config{
+		ServerMode:           serverMode,
 		BindAddr:             bindAddr,
 		Port:                 port,
+		ListenAddr:           fmt.Sprintf("%s:%d", bindAddr, port),
 		LogLevel:             loglevel,
 		PayloadUA:            payloadUA,
 		UAPattern:            uaPattern,

src/internal/log/log.go

@@ -63,9 +63,9 @@ func SetLogConf(level string) {
 	}
 }
 
-func LogHeader(version string, addr string, cfg *config.Config) {
+func LogHeader(version string, cfg *config.Config) {
 	logrus.Info("UA3F v" + version)
-	logrus.Infof("Listen on %s", addr)
+	logrus.Infof("Listen on %s", cfg.ListenAddr)
 	logrus.Infof("User-Agent: %s", cfg.PayloadUA)
 	logrus.Infof("User-Agent Regex Pattern: '%s'", cfg.UAPattern)
 	logrus.Infof("Enable Partial Replace: %v", cfg.EnablePartialReplace)

src/internal/server/server.go

@@ -0,0 +1,22 @@
+package server
+
+import (
+	"fmt"
+
+	"github.com/sunbk201/ua3f/internal/config"
+	"github.com/sunbk201/ua3f/internal/rewrite"
+	"github.com/sunbk201/ua3f/internal/server/socks5"
+)
+
+type Server interface {
+	Start() error
+}
+
+func NewServer(cfg *config.Config, rw *rewrite.Rewriter) (Server, error) {
+	switch cfg.ServerMode {
+	case config.ServerModeSocks5:
+		return socks5.New(cfg, rw), nil
+	default:
+		return nil, fmt.Errorf("unknown server mode: %s", cfg.ServerMode)
+	}
+}

src/internal/server/socks5/socks5.go

@@ -34,24 +34,22 @@ var (
 
 // Server is a minimal SOCKS5 server that delegates HTTP UA rewriting to Rewriter.
 type Server struct {
-	cfg        *config.Config
-	rw         *rewrite.Rewriter
-	listener   net.Listener
-	ListenAddr string
+	cfg      *config.Config
+	rw       *rewrite.Rewriter
+	listener net.Listener
 }
 
 // New returns a new Server with given config, rewriter, and version string.
 func New(cfg *config.Config, rw *rewrite.Rewriter) *Server {
 	return &Server{
-		cfg:        cfg,
-		rw:         rw,
-		ListenAddr: fmt.Sprintf("%s:%d", cfg.BindAddr, cfg.Port),
+		cfg: cfg,
+		rw:  rw,
 	}
 }
 
 // Start begins listening for SOCKS5 clients.
 func (s *Server) Start() (err error) {
-	if s.listener, err = net.Listen("tcp", s.ListenAddr); err != nil {
+	if s.listener, err = net.Listen("tcp", s.cfg.ListenAddr); err != nil {
 		return fmt.Errorf("listen failed: %w", err)
 	}
 

src/main.go

@@ -8,7 +8,7 @@ import (
 	"github.com/sunbk201/ua3f/internal/config"
 	"github.com/sunbk201/ua3f/internal/log"
 	"github.com/sunbk201/ua3f/internal/rewrite"
-	"github.com/sunbk201/ua3f/internal/server/socks5"
+	"github.com/sunbk201/ua3f/internal/server"
 )
 
 const version = "0.9.0"
@@ -26,8 +26,13 @@ func main() {
 		logrus.Fatal(err)
 	}
 
-	srv := socks5.New(cfg, rw)
-	log.LogHeader(version, srv.ListenAddr, cfg)
+	srv, err := server.NewServer(cfg, rw)
+	if err != nil {
+		logrus.Fatal(err)
+	}
+
+	log.LogHeader(version, cfg)
+
 	if err := srv.Start(); err != nil {
 		logrus.Fatal(err)
 	}