From 89ccbd9218e8a21952e6544e8a03fd9d86cb913d Mon Sep 17 00:00:00 2001
From: esimov <esimovaq@yahoo.com>
Date: Tue, 3 Nov 2020 09:52:56 +0200
Subject: [PATCH] Fixed flags parsing issue

---
 ws-tcp-relay.go | 99 +++++++++++++++++++++++++++----------------------
 1 file changed, 55 insertions(+), 44 deletions(-)

diff --git a/ws-tcp-relay.go b/ws-tcp-relay.go
index ff0e3ae..8f4632c 100644
--- a/ws-tcp-relay.go
+++ b/ws-tcp-relay.go
@@ -12,8 +12,61 @@ import (
 	"golang.org/x/net/websocket"
 )
 
-var tcpAddress string
-var binaryMode bool
+var (
+	tcpAddress string
+	certFile   string
+	keyFile    string
+	port       uint
+	binaryMode bool
+	fs         flag.FlagSet
+)
+
+func init() {
+	fs = *flag.NewFlagSet("Commands", flag.ExitOnError)
+	fs.UintVar(&port, "p", 4223, "The port to listen on")
+	fs.UintVar(&port, "port", 4223, "The port to listen on")
+	fs.StringVar(&certFile, "tlscert", "", "TLS cert file path")
+	fs.StringVar(&keyFile, "tlskey", "", "TLS key file path")
+	fs.BoolVar(&binaryMode, "b", false, "Use binary frames instead of text frames")
+	fs.BoolVar(&binaryMode, "binary", false, "Use binary frames instead of text frames")
+
+	fs.Usage = func() {
+		fmt.Fprintf(os.Stderr, "Usage: %s <tcpTargetAddress>\n", os.Args[0])
+		fs.PrintDefaults()
+	}
+
+	if len(os.Args) <= 1 {
+		fmt.Fprintln(os.Stderr, "No address specified")
+		os.Exit(1)
+	}
+
+	if os.Args[1] == "--help" || os.Args[1] == "-h" {
+		fmt.Fprintf(os.Stderr, "Usage: %s <tcpTargetAddress>\n", os.Args[0])
+		fs.PrintDefaults()
+		os.Exit(0)
+	}
+	fs.Parse(os.Args[2:])
+}
+
+func main() {
+	portString := fmt.Sprintf(":%d", port)
+
+	tcpAddress = os.Args[1]
+
+	log.Printf("[INFO] Listening on %s\n", portString)
+
+	http.Handle("/", websocket.Handler(relayHandler))
+
+	var err error
+	if certFile != "" && keyFile != "" {
+		err = http.ListenAndServeTLS(portString, certFile, keyFile, nil)
+	} else {
+		err = http.ListenAndServe(portString, nil)
+	}
+	if err != nil {
+		log.Fatal(err)
+	}
+}
 
 func copyWorker(dst io.Writer, src io.Reader, doneCh chan<- bool) {
 	io.Copy(dst, src)
@@ -41,45 +94,3 @@ func relayHandler(ws *websocket.Conn) {
 	ws.Close()
 	<-doneCh
 }
-
-func usage() {
-	fmt.Fprintf(os.Stderr, "Usage: %s <tcpTargetAddress>\n", os.Args[0])
-	flag.PrintDefaults()
-}
-
-func main() {
-	var port uint
-	var certFile string
-	var keyFile string
-
-	flag.UintVar(&port, "p", 4223, "The port to listen on")
-	flag.UintVar(&port, "port", 4223, "The port to listen on")
-	flag.StringVar(&certFile, "tlscert", "", "TLS cert file path")
-	flag.StringVar(&keyFile, "tlskey", "", "TLS key file path")
-	flag.BoolVar(&binaryMode, "b", false, "Use binary frames instead of text frames")
-	flag.BoolVar(&binaryMode, "binary", false, "Use binary frames instead of text frames")
-	flag.Usage = usage
-	flag.Parse()
-
-	tcpAddress = flag.Arg(0)
-	if tcpAddress == "" {
-		fmt.Fprintln(os.Stderr, "No address specified")
-		os.Exit(1)
-	}
-
-	portString := fmt.Sprintf(":%d", port)
-
-	log.Printf("[INFO] Listening on %s\n", portString)
-
-	http.Handle("/", websocket.Handler(relayHandler))
-
-	var err error
-	if certFile != "" && keyFile != "" {
-		err = http.ListenAndServeTLS(portString, certFile, keyFile, nil)
-	} else {
-		err = http.ListenAndServe(portString, nil)
-	}
-	if err != nil {
-		log.Fatal(err)
-	}
-}