From 3778160369a859a16ea4995956f09f981354457d Mon Sep 17 00:00:00 2001
From: thebinary <binary4bytes@gmail.com>
Date: Tue, 28 Jul 2020 20:56:07 +0545
Subject: [PATCH] add -to support for port only

---
 main.go | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/main.go b/main.go
index d97d75b..cd86460 100644
--- a/main.go
+++ b/main.go
@@ -2,11 +2,14 @@ package main
 
 import (
 	"flag"
+	"fmt"
 	"log"
+	"net"
 	"os"
 	"os/signal"
 	"os/user"
 	"path"
+	"strconv"
 	"strings"
 )
 
@@ -22,6 +25,23 @@ var logReqHeader bool
 var logRespHeader bool
 var logRespBody bool
 
+const flagToDescr = `local service adddress
+supported address formats:
+- [port]
+- :[port]
+- [ip]:[port]
+Note: [port] format (without colon) signifies 127.0.0.1:8080
+`
+
+func sanitizeLocalAddress(addr string) (serviceAddr string, err error) {
+	// if address has only digits consider it as port
+	if port, err := strconv.Atoi(addr); err == nil {
+		return fmt.Sprintf("127.0.0.1:%d", port), nil
+	}
+	_, _, err = net.SplitHostPort(addr)
+	return addr, err
+}
+
 func init() {
 	// get current system user info
 	cuUser, err := user.Current()
@@ -33,7 +53,7 @@ func init() {
 	flag.StringVar(&porter, "using", "http://localhost:8888", "porter server address; API endpoint address")
 	flag.StringVar(&serviceName, "as", "", "service name; defaults to what is given by porter server")
 	flag.StringVar(&tunnelAddr, "via", "", "porter tunnel endpoint adddress; SSH address; defaults to 'using_host:22'")
-	flag.StringVar(&localAddr, "to", "127.0.0.1:8080", "local listener adddress; local service listening address")
+	flag.StringVar(&localAddr, "to", "127.0.0.1:8080", flagToDescr)
 	flag.StringVar(&userName, "with-user", cuUser.Username, "username to use for porter")
 	flag.StringVar(&keyFile, "with-key", defKeyFile, "private key to use for porter")
 	//TODO: this remoteAddr should be handled by porter
@@ -43,6 +63,12 @@ func init() {
 	flag.BoolVar(&logRespBody, "log.response.body", false, "log request body")
 	flag.Parse()
 
+	localAddr, err = sanitizeLocalAddress(localAddr)
+	if err != nil {
+		fmt.Errorf("invalid form of -to address: ")
+		os.Exit(1)
+	}
+
 	// extract default tunnelAddr from porter url
 	// eg: http://ported.example.com:8888 -> extract and convert -> ported.example.com:22
 	if tunnelAddr == "" {