From 89c25ad79d7658d3de52ec582e77e23e34aa159f Mon Sep 17 00:00:00 2001
From: Ivan Daniluk <ivan.daniluk@gmail.com>
Date: Sun, 3 May 2015 20:42:21 +0300
Subject: [PATCH] Fixed ports parsing

---
 utils.go      | 42 ++++++++++++++++++++++++++++++++++++++++++
 utils_test.go | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+)

diff --git a/utils.go b/utils.go
index 196c6ee..cfc158f 100644
--- a/utils.go
+++ b/utils.go
@@ -3,6 +3,7 @@ package main
 import (
 	"errors"
 	"fmt"
+	"net"
 	"path/filepath"
 	"strings"
 
@@ -36,6 +37,47 @@ func BaseCommand(cmdline []string) string {
 
 // ParsePorts converts comma-separated ports into strings slice
 func ParsePorts(s string) ([]string, error) {
+	var (
+		ports []string
+		err   error
+	)
+	// Try simple mode, ports only ("1234-1235,80")
+	ports, err = parseRange(s)
+	if err == nil {
+		return ports, nil
+	}
+
+	var ErrParsePorts = fmt.Errorf("cannot parse ports argument")
+
+	// else, try host:ports notation ("localhost:1234-1235,remote:2000,2345")
+	fields := strings.FieldsFunc(s, func(r rune) bool { return r == ',' })
+	for _, field := range fields {
+		// split host:ports
+		var host, portsRange string
+		parts := strings.FieldsFunc(field, func(r rune) bool { return r == ':' })
+		if len(parts) == 1 {
+			host = "localhost"
+		} else if len(parts) == 2 {
+			host, portsRange = parts[0], parts[1]
+		} else {
+			return nil, ErrParsePorts
+		}
+
+		pp, err := parseRange(portsRange)
+		if err != nil {
+			return nil, ErrParsePorts
+		}
+
+		for _, p := range pp {
+			addr := net.JoinHostPort(host, p)
+			ports = append(ports, addr)
+		}
+	}
+
+	return ports, nil
+}
+
+func parseRange(s string) ([]string, error) {
 	portsInt, err := ranges.Parse(s)
 	if err != nil {
 		return nil, err
diff --git a/utils_test.go b/utils_test.go
index dc1b4c1..492c684 100644
--- a/utils_test.go
+++ b/utils_test.go
@@ -25,3 +25,44 @@ func TestUtils(t *testing.T) {
 		t.Fatalf("vars should contain 4 elements, but has %d", len(vars))
 	}
 }
+
+func TestPorts(t *testing.T) {
+	arg := "1234,1235"
+	ports, err := ParsePorts(arg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(ports) != 2 || ports[0] != "1234" {
+		t.Fatalf("ParsePorts returns wrong data: %v", ports)
+	}
+
+	arg = "1234-1237,2000"
+	ports, err = ParsePorts(arg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(ports) != 5 || ports[0] != "1234" || ports[4] != "2000" {
+		t.Fatalf("ParsePorts returns wrong data: %v", ports)
+	}
+
+	arg = "localhost:2000-2002,remote:1234-1235"
+	ports, err = ParsePorts(arg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(ports) != 5 || ports[0] != "localhost:2000" || ports[4] != "remote:1235" {
+		t.Fatalf("ParsePorts returns wrong data: %v", ports)
+	}
+
+	arg = "localhost:2000-2002,remote:1234-1235,some:weird:1234-123input"
+	_, err = ParsePorts(arg)
+	if err == nil {
+		t.Fatalf("err shouldn't be nil")
+	}
+
+	arg = "string,sdasd"
+	_, err = ParsePorts(arg)
+	if err == nil {
+		t.Fatalf("err shouldn't be nil")
+	}
+}