driverutil: refactor splitPortProto to parse port as str

andrestc · andrestc · commit 16bf6769edcd · 2016-08-22T09:29:51.000-03:00
Signed-off-by: André Carvalho &lt;andre.carvalho@corp.globo.com&gt;
diff --git a/drivers/amazonec2/amazonec2.go b/drivers/amazonec2/amazonec2.go
@@ -1108,10 +1108,14 @@ func (d *Driver) configureSecurityGroupPermissions(group *ec2.SecurityGroup) ([]
 		if err != nil {
 			return nil, err
 		}
+		portNum, err := strconv.Atoi(port)
+		if err != nil {
+			return nil, fmt.Errorf("invalid port number %s: %s", port, err)
+		}
 		perms = append(perms, &ec2.IpPermission{
 			IpProtocol: aws.String(protocol),
-			FromPort:   aws.Int64(int64(port)),
-			ToPort:     aws.Int64(int64(port)),
+			FromPort:   aws.Int64(int64(portNum)),
+			ToPort:     aws.Int64(int64(portNum)),
 			IpRanges:   []*ec2.IpRange{{CidrIp: aws.String(ipRange)}},
 		})
 	}
diff --git a/drivers/azure/util.go b/drivers/azure/util.go
@@ -5,7 +5,6 @@ import (
 	"io/ioutil"
 	"net"
 	"net/url"
-	"strconv"
 	"strings"
 
 	"github.com/Azure/azure-sdk-for-go/arm/network"
@@ -128,11 +127,10 @@ func (d *Driver) getSecurityRules(extraPorts []string) (*[]network.SecurityRule,
 	// extra port numbers requested by user
 	basePri := 1000
 	for i, p := range extraPorts {
-		n, protocol, err := driverutil.SplitPortProto(p)
+		port, protocol, err := driverutil.SplitPortProto(p)
 		if err != nil {
 			return nil, err
 		}
-		port := strconv.Itoa(n)
 		proto, err := parseSecurityRuleProtocol(protocol)
 		if err != nil {
 			return nil, err
diff --git a/drivers/driverutil/util.go b/drivers/driverutil/util.go
@@ -1,23 +1,16 @@
 package driverutil
 
-import (
-	"fmt"
-	"strconv"
-	"strings"
-)
+import "strings"
 
 // SplitPortProto splits a string in the format port/protocol, defaulting
 // protocol to "tcp" if not provided.
-func SplitPortProto(raw string) (port int, protocol string, err error) {
+func SplitPortProto(raw string) (port string, protocol string, err error) {
 	parts := strings.Split(raw, "/")
 	if len(parts) == 1 {
 		protocol = "tcp"
 	} else {
 		protocol = parts[1]
 	}
-	port, err = strconv.Atoi(parts[0])
-	if err != nil {
-		return 0, "", fmt.Errorf("invalid port number %s: %s", parts[0], err)
-	}
+	port = parts[0]
 	return port, protocol, nil
 }
diff --git a/drivers/driverutil/util_test.go b/drivers/driverutil/util_test.go
@@ -9,14 +9,13 @@ import (
 func TestSplitPortProtocol(t *testing.T) {
 	tests := []struct {
 		raw           string
-		expectedPort  int
+		expectedPort  string
 		expectedProto string
 		expectedErr   bool
 	}{
-		{"8080/tcp", 8080, "tcp", false},
-		{"90/udp", 90, "udp", false},
-		{"80", 80, "tcp", false},
-		{"abc", 0, "", true},
+		{"8080/tcp", "8080", "tcp", false},
+		{"90/udp", "90", "udp", false},
+		{"80", "80", "tcp", false},
 	}
 
 	for _, tc := range tests {
