[main] Add -- cmd flag separator to start different Gost instances

cmd/gost/main.go

@@ -7,6 +7,8 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"sync"
+	"strings"
 	"runtime"
 
 	_ "net/http/pprof"
@@ -16,86 +18,92 @@ import (
 )
 
 var (
-	configureFile string
-	baseCfg       = &baseConfig{}
-	pprofAddr     string
 	pprofEnabled  = os.Getenv("PROFILING") != ""
 )
 
 func init() {
 	gost.SetLogger(&gost.LogLogger{})
 
-	var (
+	// TODO - Generate different certificates for each worker
-		printVersion bool
+	generateTLSCertificate()
-	)
-
-	flag.Var(&baseCfg.route.ChainNodes, "F", "forward address, can make a forward chain")
-	flag.Var(&baseCfg.route.ServeNodes, "L", "listen address, can listen on multiple ports (required)")
-	flag.IntVar(&baseCfg.route.Mark, "M", 0, "Specify out connection mark")
-	flag.StringVar(&configureFile, "C", "", "configure file")
-	flag.StringVar(&baseCfg.route.Interface, "I", "", "Interface to bind")
-	flag.BoolVar(&baseCfg.Debug, "D", false, "enable debug log")
-	flag.BoolVar(&printVersion, "V", false, "print version")
-	if pprofEnabled {
-		flag.StringVar(&pprofAddr, "P", ":6060", "profiling HTTP server address")
-	}
-	flag.Parse()
-
-	if printVersion {
-		fmt.Fprintf(os.Stdout, "gost %s (%s %s/%s)\n",
-			gost.Version, runtime.Version(), runtime.GOOS, runtime.GOARCH)
-		os.Exit(0)
-	}
-
-	if configureFile != "" {
-		_, err := parseBaseConfig(configureFile)
-		if err != nil {
-			log.Log(err)
-			os.Exit(1)
-		}
-	}
-	if flag.NFlag() == 0 {
-		flag.PrintDefaults()
-		os.Exit(0)
-	}
 }
 
 func main() {
-	if pprofEnabled {
+	var wg sync.WaitGroup
-		go func() {
+	wg.Add(1)  // Gost must exit if any of the workers exit
-			log.Log("profiling server on", pprofAddr)
-			log.Log(http.ListenAndServe(pprofAddr, nil))
-		}()
-	}
 
-	// NOTE: as of 2.6, you can use custom cert/key files to initialize the default certificate.
+	// Split os.Args using -- and create a worker with each slice
-	tlsConfig, err := tlsConfig(defaultCertFile, defaultKeyFile, "")
+	args := strings.Split(" " + strings.Join(os.Args[1:], "  ") + " ", " -- ")
-	if err != nil {
+	if strings.Join(args, "") == "" {
-		// generate random self-signed certificate.
+		// Fix to show gost help if the resulting array is empty
-		cert, err := gost.GenCertificate()
+		args[0] = " "
-		if err != nil {
-			log.Log(err)
-			os.Exit(1)
 	}
-		tlsConfig = &tls.Config{
+	for wid, wargs := range args {
-			Certificates: []tls.Certificate{cert},
+		if wargs != "" {
+			go worker(wid, wargs, &wg)
 		}
-	} else {
-		log.Log("load TLS certificate files OK")
 	}
-
+	wg.Wait()
-	gost.DefaultTLSConfig = tlsConfig
-
-	if err := start(); err != nil {
-		log.Log(err)
-		os.Exit(1)
-	}
-
-	select {}
 }
 
-func start() error {
+func worker(id int, args string, wg *sync.WaitGroup) {
+	defer wg.Done()
+
+	var (
+		configureFile string
+		baseCfg       = &baseConfig{}
+		pprofAddr     string
+	)
+
+	init := func () error {
+		var printVersion bool
+
+		wf := flag.NewFlagSet(os.Args[0], flag.ExitOnError)
+
+		wf.Var(&baseCfg.route.ChainNodes, "F", "forward address, can make a forward chain")
+		wf.Var(&baseCfg.route.ServeNodes, "L", "listen address, can listen on multiple ports (required)")
+		wf.StringVar(&configureFile, "C", "", "configure file")
+		wf.BoolVar(&baseCfg.Debug, "D", false, "enable debug log")
+		wf.BoolVar(&printVersion, "V", false, "print version")
+
+		if pprofEnabled {
+			// Every worker uses a different profiling server by default
+			wf.StringVar(&pprofAddr, "P", fmt.Sprintf(":606%d", id), "profiling HTTP server address")
+		}
+
+		wf.Parse(strings.Fields(args))
+
+		if printVersion {
+			fmt.Fprintf(os.Stdout, "gost %s (%s %s/%s)\n", gost.Version, runtime.Version(), runtime.GOOS, runtime.GOARCH)
+			os.Exit(0)
+		} else if wf.NFlag() == 0 {
+			wf.Usage()
+			os.Exit(0)
+		} else if configureFile != "" {
+			err := parseBaseConfig(configureFile, baseCfg)
+			if err != nil {
+				return err
+			}
+		}
+
+		if baseCfg.route.ServeNodes.String() == "[]" {
+			configErrMsg := ""
+			if configureFile != "" {
+				configErrMsg = " or ServeNodes inside config file (-C)"
+			}
+			fmt.Fprintf(os.Stderr, "\n[!] Error: Missing -L flag%s\n\n", configErrMsg)
+			wf.Usage()
+			os.Exit(1)
+		}
+
+		return nil
+	}
+
+	start := func () error {
+		// TODO - Make debug worker independent
+		if ! gost.Debug {
 			gost.Debug = baseCfg.Debug
+		}
 
 		var routers []router
 		rts, err := baseCfg.route.GenRouters()
@@ -120,4 +128,52 @@ func start() error {
 		}
 
 		return nil
+	}
+
+	main := func () error {
+		if pprofEnabled {
+			go func() {
+				log.Log("profiling server on", pprofAddr)
+				log.Log(http.ListenAndServe(pprofAddr, nil))
+			}()
+		}
+
+		err := start()
+		return err
+	}
+
+	if err := init(); err != nil {
+		log.Log(err)
+		return
+	}
+	if err := main(); err != nil {
+		log.Log(err)
+		return
+	}
+
+	// Allow local functions to be garbage-collected
+	init = nil
+	main = nil
+	start = nil
+
+	select {}
+}
+
+func generateTLSCertificate() {
+	// NOTE: as of 2.6, you can use custom cert/key files to initialize the default certificate.
+	tlsConfig, err := tlsConfig(defaultCertFile, defaultKeyFile, "")
+	if err != nil {
+		// generate random self-signed certificate.
+		cert, err := gost.GenCertificate()
+		if err != nil {
+			log.Log(err)
+			os.Exit(1)
+		}
+		tlsConfig = &tls.Config{
+			Certificates: []tls.Certificate{cert},
+		}
+	} else {
+		log.Log("load TLS certificate files OK")
+	}
+	gost.DefaultTLSConfig = tlsConfig
 }