azure: add support for custom protocol in --azure-open-port

This makes the --azure-open-port flag behave like
--amazonec2-open-port flag, enabling the user to provide a
protocol for the port, in the format port/protocol.

Signed-off-by: André Carvalho <andre.carvalho@corp.globo.com>

drivers/azure/util.go

@@ -5,6 +5,7 @@ import (
 	"io/ioutil"
 	"net"
 	"net/url"
+	"strconv"
 	"strings"
 
 	"github.com/Azure/azure-sdk-for-go/arm/network"
@@ -12,6 +13,7 @@ import (
 	"github.com/Azure/go-autorest/autorest/to"
 	"github.com/docker/machine/drivers/azure/azureutil"
 	"github.com/docker/machine/drivers/azure/logutil"
+	"github.com/docker/machine/drivers/util"
 	"github.com/docker/machine/libmachine/log"
 	"github.com/docker/machine/libmachine/ssh"
 	"github.com/docker/machine/libmachine/state"
@@ -81,7 +83,7 @@ func (d *Driver) generateSSHKey(ctx *azureutil.DeploymentContext) error {
 // getSecurityRules creates network security group rules based on driver
 // configuration such as SSH port, docker port and swarm port.
 func (d *Driver) getSecurityRules(extraPorts []string) (*[]network.SecurityRule, error) {
-	mkRule := func(priority int, name, description, srcPort, dstPort string) network.SecurityRule {
+	mkRule := func(priority int, name, description, srcPort, dstPort string, proto network.SecurityRuleProtocol) network.SecurityRule {
 		return network.SecurityRule{
 			Name: to.StringPtr(name),
 			Properties: &network.SecurityRulePropertiesFormat{
@@ -92,7 +94,7 @@ func (d *Driver) getSecurityRules(extraPorts []string) (*[]network.SecurityRule,
 				DestinationPortRange:     to.StringPtr(dstPort),
 				Access:                   network.Allow,
 				Direction:                network.Inbound,
-				Protocol:                 network.TCP,
+				Protocol:                 proto,
 				Priority:                 to.Int32Ptr(int32(priority)),
 			},
 		}
@@ -102,8 +104,8 @@ func (d *Driver) getSecurityRules(extraPorts []string) (*[]network.SecurityRule,
 
 	// Base ports to be opened for any machine
 	rl := []network.SecurityRule{
-		mkRule(100, "SSHAllowAny", "Allow ssh from public Internet", "*", fmt.Sprintf("%d", d.BaseDriver.SSHPort)),
-		mkRule(300, "DockerAllowAny", "Allow docker engine access (TLS-protected)", "*", fmt.Sprintf("%d", d.DockerPort)),
+		mkRule(100, "SSHAllowAny", "Allow ssh from public Internet", "*", fmt.Sprintf("%d", d.BaseDriver.SSHPort), network.TCP),
+		mkRule(300, "DockerAllowAny", "Allow docker engine access (TLS-protected)", "*", fmt.Sprintf("%d", d.DockerPort), network.TCP),
 	}
 
 	// Open swarm port if configured
@@ -118,16 +120,25 @@ func (d *Driver) getSecurityRules(extraPorts []string) (*[]network.SecurityRule,
 		if err != nil {
 			return nil, fmt.Errorf("Could not parse swarm port in %q: %v", u.Host, err)
 		}
-		rl = append(rl, mkRule(500, "DockerSwarmAllowAny", "Allow swarm manager access (TLS-protected)", "*", swarmPort))
+		rl = append(rl, mkRule(500, "DockerSwarmAllowAny", "Allow swarm manager access (TLS-protected)", "*", swarmPort, network.TCP))
 	} else {
 		log.Debug("Swarm host is not configured.")
 	}
 
 	// extra port numbers requested by user
 	basePri := 1000
-	for i, port := range extraPorts {
-		log.Debugf("User-requested port number to be opened on NSG: %v", port)
-		r := mkRule(basePri+i, fmt.Sprintf("Port%sAllowAny", port), "User requested port to be accessible from Internet via docker-machine", "*", port)
+	for i, p := range extraPorts {
+		n, protocol, err := util.SplitPortProto(p)
+		if err != nil {
+			return nil, err
+		}
+		port := strconv.Itoa(n)
+		proto, err := parseSecurityRuleProtocol(protocol)
+		if err != nil {
+			return nil, err
+		}
+		log.Debugf("User-requested port to be opened on NSG: %v/%s", port, proto)
+		r := mkRule(basePri+i, fmt.Sprintf("Port%s%sAllowAny", port, proto), "User requested port to be accessible from Internet via docker-machine", "*", port, proto)
 		rl = append(rl, r)
 	}
 	log.Debugf("Total NSG rules: %d", len(rl))
@@ -193,3 +204,18 @@ func parseVirtualNetwork(name string, defaultRG string) (string, string) {
 	}
 	return defaultRG, name
 }
+
+// parseSecurityRuleProtocol parses a protocol string into a network.SecurityRuleProtocol
+// and returns error if the protocol is not supported
+func parseSecurityRuleProtocol(proto string) (network.SecurityRuleProtocol, error) {
+	switch strings.ToLower(proto) {
+	case "tcp":
+		return network.TCP, nil
+	case "udp":
+		return network.UDP, nil
+	case "*":
+		return network.Asterisk, nil
+	default:
+		return "", fmt.Errorf("invalid protocol %s", proto)
+	}
+}

drivers/azure/util_test.go

@@ -0,0 +1,31 @@
+package azure
+
+import (
+	"testing"
+
+	"github.com/Azure/azure-sdk-for-go/arm/network"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestParseSecurityRuleProtocol(t *testing.T) {
+	tests := []struct {
+		raw           string
+		expectedProto network.SecurityRuleProtocol
+		expectedErr   bool
+	}{
+		{"tcp", network.TCP, false},
+		{"udp", network.UDP, false},
+		{"*", network.Asterisk, false},
+		{"Invalid", "", true},
+	}
+
+	for _, tc := range tests {
+		proto, err := parseSecurityRuleProtocol(tc.raw)
+		assert.Equal(t, tc.expectedProto, proto)
+		if tc.expectedErr {
+			assert.Error(t, err)
+		} else {
+			assert.NoError(t, err)
+		}
+	}
+}