From 436a39b1d6a89a79986606bb6f6391d438642e3c Mon Sep 17 00:00:00 2001
From: Kai Luo <gluokai@gmail.com>
Date: Mon, 25 Dec 2023 10:04:18 +0800
Subject: [PATCH] Refactor options

---
 i3/main.go           | 36 ++++++++++--------------------------
 j2/main.go           |  1 -
 relayer/intrinsic.go | 16 ++++++++++------
 3 files changed, 20 insertions(+), 33 deletions(-)

diff --git a/i3/main.go b/i3/main.go
index e61c21f..872479d 100644
--- a/i3/main.go
+++ b/i3/main.go
@@ -14,7 +14,6 @@ import (
 	"math/rand"
 	"net"
 	"strings"
-	"sync"
 
 	"github.com/bzEq/bx/core"
 	"github.com/bzEq/bx/relayer"
@@ -22,38 +21,24 @@ import (
 
 var options struct {
 	Local          string
+	LocalUDP       string
 	LocalHTTPProxy string
 	Next           string
 	ProtocolName   string
-	NumConn        int
 	UseTLS         bool
-	NoUDP          bool
 }
 
-func startRelayers() {
-	addrs := strings.Split(options.Local, ",")
-	var wg sync.WaitGroup
-	for _, addr := range addrs {
-		wg.Add(1)
-		go func(addr string) {
-			defer wg.Done()
-			startRelayer(addr)
-		}(addr)
-	}
-	wg.Wait()
-}
-
-func startRelayer(localAddr string) {
+func startRelayer() {
 	r := &relayer.IntrinsicRelayer{}
-	r.Local = localAddr
+	r.Local = options.Local
+	r.LocalUDP = options.LocalUDP
+	r.NumUDPMux = 4
 	r.LocalHTTPProxy = options.LocalHTTPProxy
 	r.RelayProtocol = options.ProtocolName
-	r.NumUDPMux = options.NumConn
-	r.NoUDP = options.NoUDP
 	if options.Next != "" {
 		r.Next = strings.Split(options.Next, ",")
 	}
-	if options.UseTLS && len(r.Next) != 0 {
+	if options.UseTLS && !r.IsEndPoint() {
 		config := &tls.Config{InsecureSkipVerify: true, NextProtos: []string{options.ProtocolName}}
 		r.Dial = func(network, address string) (net.Conn, error) {
 			return tls.Dial(network, address, config)
@@ -63,7 +48,7 @@ func startRelayer(localAddr string) {
 			return net.Dial(network, address)
 		}
 	}
-	if options.UseTLS && len(r.Next) == 0 {
+	if options.UseTLS && r.IsEndPoint() {
 		config, err := core.CreateBarebonesTLSConfig(options.ProtocolName)
 		if err != nil {
 			log.Println(err)
@@ -86,17 +71,16 @@ func main() {
 	rand.Seed(seed)
 	var debug bool
 	flag.StringVar(&options.Local, "l", "localhost:1080", "Listen address of this relayer")
+	flag.StringVar(&options.LocalUDP, "u", "", "UDP listen address of this relayer")
 	flag.StringVar(&options.Next, "n", "", "Address of next-hop relayer")
 	flag.StringVar(&options.LocalHTTPProxy, "http_proxy", "", "Enable this relayer serving as http proxy")
 	flag.StringVar(&options.ProtocolName, "proto", "", "Name of relay protocol")
-	flag.BoolVar(&options.UseTLS, "use_tls", false, "Use TLS")
-	flag.IntVar(&options.NumConn, "j", 4, "Number of connections for UDP Mux")
-	flag.BoolVar(&options.NoUDP, "no_udp", false, "Disable UDP support")
+	flag.BoolVar(&options.UseTLS, "tls", false, "Wrap traffic in TLS")
 	flag.BoolVar(&debug, "debug", false, "Enable debug logging")
 	flag.Parse()
 	if !debug {
 		log.SetOutput(ioutil.Discard)
 	}
 	log.SetFlags(log.LstdFlags | log.Lshortfile)
-	startRelayers()
+	startRelayer()
 }
diff --git a/j2/main.go b/j2/main.go
index 4f464d9..45b6916 100644
--- a/j2/main.go
+++ b/j2/main.go
@@ -24,7 +24,6 @@ var options struct {
 	Local        string
 	Next         string
 	ProtocolName string
-	NumConn      int
 	UseTLS       bool
 }
 
diff --git a/relayer/intrinsic.go b/relayer/intrinsic.go
index 1bd5d54..57d666e 100644
--- a/relayer/intrinsic.go
+++ b/relayer/intrinsic.go
@@ -18,12 +18,12 @@ import (
 type IntrinsicRelayer struct {
 	Listen         func(string, string) (net.Listener, error)
 	Local          string
+	LocalUDP       string
+	NumUDPMux      int
 	LocalHTTPProxy string
 	Dial           func(string, string) (net.Conn, error)
 	Next           []string
 	RelayProtocol  string
-	NumUDPMux      int
-	NoUDP          bool
 	udpAddr        *net.UDPAddr
 }
 
@@ -58,7 +58,7 @@ func (self *IntrinsicRelayer) startLocalHTTPProxy() error {
 }
 
 func (self *IntrinsicRelayer) startLocalUDPServer() error {
-	laddr, err := net.ResolveUDPAddr("udp", self.Local)
+	laddr, err := net.ResolveUDPAddr("udp", self.LocalUDP)
 	if err != nil {
 		return err
 	}
@@ -91,8 +91,12 @@ func (self *IntrinsicRelayer) startLocalUDPServer() error {
 	return nil
 }
 
+func (self *IntrinsicRelayer) IsEndPoint() bool {
+	return len(self.Next) == 0
+}
+
 func (self *IntrinsicRelayer) Run() {
-	if len(self.Next) != 0 && !self.NoUDP {
+	if !self.IsEndPoint() && self.LocalUDP != "" {
 		if err := self.startLocalUDPServer(); err != nil {
 			log.Println(err)
 			return
@@ -105,7 +109,7 @@ func (self *IntrinsicRelayer) Run() {
 	}
 	defer ln.Close()
 	// self.LocalHTTPProxy relies on socks proxy.
-	if len(self.Next) != 0 && self.LocalHTTPProxy != "" {
+	if !self.IsEndPoint() && self.LocalHTTPProxy != "" {
 		if err := self.startLocalHTTPProxy(); err != nil {
 			log.Println(err)
 			return
@@ -117,7 +121,7 @@ func (self *IntrinsicRelayer) Run() {
 			log.Println(err)
 			break
 		}
-		if len(self.Next) == 0 {
+		if self.IsEndPoint() {
 			go self.ServeAsEndRelayer(c)
 		} else {
 			go self.ServeAsLocalRelayer(c)