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,6 +23,7 @@ type Config struct {
 
 func Parse() (*Config, bool) {
 	var (
+		serverMode string
 		bindAddr   string
 		port       int
 		loglevel   string
@@ -22,6 +33,7 @@ func Parse() (*Config, 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

@@ -37,7 +37,6 @@ type Server struct {
 	cfg      *config.Config
 	rw       *rewrite.Rewriter
 	listener net.Listener
-	ListenAddr string
 }
 
 // New returns a new Server with given config, rewriter, and version string.
@@ -45,13 +44,12 @@ func New(cfg *config.Config, rw *rewrite.Rewriter) *Server {
 	return &Server{
 		cfg: cfg,
 		rw:  rw,
-		ListenAddr: fmt.Sprintf("%s:%d", cfg.BindAddr, cfg.Port),
 	}
 }
 
 // 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)
+	srv, err := server.NewServer(cfg, rw)
-	log.LogHeader(version, srv.ListenAddr, cfg)
+	if err != nil {
+		logrus.Fatal(err)
+	}
+
+	log.LogHeader(version, cfg)
+
 	if err := srv.Start(); err != nil {
 		logrus.Fatal(err)
 	}