Support protocol argument

README.md

@@ -14,10 +14,11 @@ Wait until an address become available.
 
 ### Options
 
-- `-address`: Address (e.g. http://google.com or tcp://mysql_ip:mysql_port) - *former **full-connection***
-- `-status`: Expected status that address should return (e.g. 200)
+- `-address`: Address (e.g. http://google.com, tcp://mysql-ip:port, ssh://ip:port) - *former **full-connection***
+- `-proto`: Protocol to use during the connection
 - `-host`: Host to connect
 - `-port`: Port to connect (default 80)
+- `-status`: Expected status that address should return (e.g. 200)
 - `-timeout`: Seconds to wait until the address become available
 - `-retry`: Milliseconds to wait between retries (default 500)
 - `-debug`: Enable debug

connection.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net"
 	"net/url"
+	"regexp"
 	"strconv"
 )
 
@@ -14,13 +15,23 @@ type Connection struct {
 	URL         *url.URL
 }
 
+var defaultProtPorts = map[string]string{
+	"http":  "80",
+	"https": "443",
+	"ssh":   "22",
+}
+
 // BuildConn build a connection structure.
 // This connection data can later be used as a common structure
 // by the functions that will check if the target is available.
 func BuildConn(cfg *Config) (*Connection, error) { // nolint gocyclo
 	address := cfg.Address
-	if cfg.Host != "" {
+	if address == "" && cfg.Host != "" {
 		address = net.JoinHostPort(cfg.Host, strconv.Itoa(cfg.Port))
+
+		if cfg.Protocol != "" {
+			address = fmt.Sprintf("%s://%s", cfg.Protocol, address)
+		}
 	}
 
 	if address == "" {
@@ -41,12 +52,23 @@ func BuildConn(cfg *Config) (*Connection, error) { // nolint gocyclo
 		return nil, fmt.Errorf("Error parsing connection address: %v", err)
 	}
 
-	if u.Hostname() == "" {
+	if u == nil || u.Hostname() == "" {
 		return nil, fmt.Errorf("Couldn't parse address: %s", address)
 	}
 
+	p := u.Port()
+
+	// resolve default port based on the provided scheme
+	if p == "0" {
+		if dp, ok := defaultProtPorts[u.Scheme]; ok {
+			var re = regexp.MustCompile(`:0$`)
+			u.Host = re.ReplaceAllString(u.Host, ":"+dp)
+		}
+	}
+
+	// resolve default scheme based on the provided port
 	if u.Scheme == "" {
-		if p := u.Port(); p == "80" {
+		if p == "80" {
 			u.Scheme = "http"
 		} else if p == "443" {
 			u.Scheme = "https"

connection_test.go

@@ -8,6 +8,7 @@ import (
 
 func TestBuildConn(t *testing.T) {
 	type input struct {
+		proto   string
 		host    string
 		port    int
 		address string
@@ -35,9 +36,9 @@ func TestBuildConn(t *testing.T) {
 			&expected{netType: "tcp", host: "localhost:90", address: "tcp://localhost:90"},
 		},
 		{
-			"Should ignore the address when the host is given",
+			"Should ignore protocol, host or port when an address is given",
 			input{host: "localhost", port: 90, address: "tcp://remotehost:10"},
-			&expected{netType: "tcp", host: "localhost:90", address: "tcp://localhost:90"},
+			&expected{netType: "tcp", host: "remotehost:10", address: "tcp://remotehost:10"},
 		},
 		{
 			"Should be able to craete a connection given a address",

main.go

@@ -15,13 +15,14 @@ var VERSION string
 
 // Config describes the connection config
 type Config struct {
-	Host    string            `json:"host"`
-	Port    int               `json:"port"`
-	Address string            `json:"address"`
-	Status  int               `json:"status"`
-	Timeout int               `json:"timeout"`
-	Retry   int               `json:"retry"`
-	Headers map[string]string `json:"headers"`
+	Protocol string            `json:"proto"`
+	Host     string            `json:"host"`
+	Port     int               `json:"port"`
+	Address  string            `json:"address"`
+	Status   int               `json:"status"`
+	Timeout  int               `json:"timeout"`
+	Retry    int               `json:"retry"`
+	Headers  map[string]string `json:"headers"`
 }
 
 // FileConfig describes the structure of the config json file
@@ -54,9 +55,10 @@ func main() { // nolint gocyclo
 	var fheaders arrayFlags
 
 	address := flag.String("address", "", "address (e.g. http://google.com or tcp://mysql_ip:mysql_port)")
-	status := flag.Int("status", 0, "expected status that address should return (e.g. 200")
+	proto := flag.String("proto", "", "protocol to use during the connection")
 	host := flag.String("host", "", "host to connect")
-	port := flag.Int("port", 80, "port to connect")
+	port := flag.Int("port", 0, "port to connect")
+	status := flag.Int("status", 0, "expected status that address should return (e.g. 200")
 	timeout := flag.Int("timeout", 10, "seconds to wait until the address become available")
 	retry := flag.Int("retry", 500, "milliseconds to wait between retries")
 	printVersion := flag.Bool("v", false, "show the current version")
@@ -104,13 +106,14 @@ func main() { // nolint gocyclo
 		fc = FileConfig{
 			Configs: []Config{
 				{
-					Host:    *host,
-					Port:    *port,
-					Address: *address,
-					Status:  *status,
-					Timeout: *timeout,
-					Retry:   *retry,
-					Headers: headers,
+					Protocol: *proto,
+					Host:     *host,
+					Port:     *port,
+					Address:  *address,
+					Status:   *status,
+					Timeout:  *timeout,
+					Retry:    *retry,
+					Headers:  headers,
 				},
 			},
 		}