cmd/tailcat: allow socks mode without cmd

cmd/tailcat/socks_test.go

@@ -135,3 +135,45 @@ func TestClassifySOCKSAddr(t *testing.T) {
 		})
 	}
 }
+
+func TestNormalizeListenAddrPort(t *testing.T) {
+	tests := []struct {
+		name  string
+		input string
+		want  string
+	}{
+		{
+			name:  "integer only",
+			input: "1234",
+			want:  "127.0.0.1:1234",
+		},
+		{
+			name:  "omit address",
+			input: ":1234",
+			want:  "0.0.0.0:1234",
+		},
+		{
+			name:  "omit port with IPv4 address",
+			input: "127.0.0.1",
+			want:  "127.0.0.1:0",
+		},
+		{
+			name:  "omit port with IPv6 address",
+			input: "[2001:db8::1]",
+			want:  "[2001:db8::1]:0",
+		},
+		{
+			name:  "others",
+			input: "foo",
+			want:  "foo:0",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := normalizeListenAddrPort(tt.input)
+			if got != tt.want {
+				t.Fatalf("classifyListenAddrPort(%q) = %q; want %q", tt.input, got, tt.want)
+			}
+		})
+	}
+}

cmd/tailcat/tailcat.go

@@ -117,10 +117,12 @@ as 'all_proxy' environment variable to a child process. Destination
 hostnames that are themselves address blobs are dialed as tailcat
 servers, so the <addrblob> argument is optional:
 
-	tailcat socks [<addrblob>] <cmd> [args...]
+	tailcat socks [-listen <addr:port>] [<addrblob>] [<cmd> [args...]]
 	tailcat socks <addrblob> curl http://server.tailcat:8081/
 	tailcat socks curl http://<addrblob>:8081/
 
+If you don't specify the cmd, just the proxy server will start.
+
 Parse an address blob and print its encoded fields as JSON:
 
 	tailcat parse <addrblob>
@@ -437,8 +439,32 @@ func clientMode(logf logger.Logf, connStr, optDest string) {
 	}
 }
 
+// This function will only fill in the missing part.
+// 0.0.0.0 and 0 will be used for the address and port respectively.
+// It won't validate the input, we delegate this to the following socket creation
+func normalizeListenAddrPort(s string) string {
+	if host, port, err := net.SplitHostPort(s); err == nil {
+		if host == "" {
+			host = "0.0.0.0"
+		}
+		if port == "" {
+			port = "0"
+		}
+		return net.JoinHostPort(host, port)
+	} else if port, err := strconv.ParseUint(s, 10, 16); err == nil {
+		return net.JoinHostPort("127.0.0.1", strconv.Itoa(int(port)))
+	}
+	// Assume it's hostname
+	return s + ":0"
+}
+
 func clientSOCKSMode(logf logger.Logf) {
-	args := flag.Args()[1:] // trim "socks"
+	fs := flag.NewFlagSet("socks", flag.ExitOnError)
+	listen := fs.String("listen", "127.0.0.1:0", "Proxy server's listen address")
+	fs.Parse(flag.Args()[1:]) // stripping off "socks"
+	args := fs.Args()
+
+	lisenAddrPort := normalizeListenAddrPort(*listen)
 
 	// The address blob argument is optional: destination hostnames that
 	// are themselves address blobs are dialed directly (see
@@ -458,9 +484,6 @@ func clientSOCKSMode(logf logger.Logf) {
 			}
 		}
 	}
-	if len(args) == 0 {
-		usage("tailcat socks [<addrblob>] <cmd> [args...]")
-	}
 	progArgs := args
 
 	var cl *tailcat.Client
@@ -489,7 +512,7 @@ func clientSOCKSMode(logf logger.Logf) {
 		return c
 	}
 
-	socksLn, err := net.Listen("tcp", "localhost:0")
+	socksLn, err := net.Listen("tcp", lisenAddrPort)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -512,20 +535,25 @@ func clientSOCKSMode(logf logger.Logf) {
 			return cl.DialTCP(ctx, dst.dst)
 		},
 	}
-	go func() {
-		log.Fatalf("SOCKS5 server exited: %v", ss.Serve(socksLn))
-	}()
 	socksAddr := "socks5h://" + socksLn.Addr().String()
-	logf("SOCKS running at %v", socksAddr)
-	cmd := exec.Command(progArgs[0], progArgs[1:]...)
-	cmd.Env = append(os.Environ(),
-		"all_proxy="+socksAddr,
-	)
-	cmd.Stdin = os.Stdin
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
-	if err := cmd.Run(); err != nil {
-		log.Fatal(err)
+	if len(progArgs) > 0 {
+		go func() {
+			log.Fatalf("SOCKS5 server exited: %v", ss.Serve(socksLn))
+		}()
+		logf("SOCKS running at %v", socksAddr)
+		cmd := exec.Command(progArgs[0], progArgs[1:]...)
+		cmd.Env = append(os.Environ(),
+			"all_proxy="+socksAddr,
+		)
+		cmd.Stdin = os.Stdin
+		cmd.Stdout = os.Stdout
+		cmd.Stderr = os.Stderr
+		if err := cmd.Run(); err != nil {
+			log.Fatal(err)
+		}
+	} else {
+		fmt.Printf("SOCKS running at %v\n", socksAddr)
+		log.Fatalf("SOCKS5 server exited: %v", ss.Serve(socksLn))
 	}
 }