From 85f1cb97258e080665c165fce4865334500466a1 Mon Sep 17 00:00:00 2001
From: Evan Hazlett <ejhazlett@gmail.com>
Date: Thu, 22 Jan 2015 10:06:56 -0500
Subject: [PATCH] only authorize ec2 group if needed

Signed-off-by: Evan Hazlett <ejhazlett@gmail.com>
---
 drivers/amazonec2/amazonec2.go          | 60 ++++++++++++----
 drivers/amazonec2/amazonec2_test.go     | 94 +++++++++++++++++++++++++
 drivers/amazonec2/amz/security_group.go |  8 ++-
 3 files changed, 144 insertions(+), 18 deletions(-)
 create mode 100644 drivers/amazonec2/amazonec2_test.go

diff --git a/drivers/amazonec2/amazonec2.go b/drivers/amazonec2/amazonec2.go
index 2573c69707..1ecef1d94b 100644
--- a/drivers/amazonec2/amazonec2.go
+++ b/drivers/amazonec2/amazonec2.go
@@ -28,6 +28,7 @@ const (
 	ipRange                  = "0.0.0.0/0"
 	dockerConfigDir          = "/etc/docker"
 	machineSecurityGroupName = "docker-machine"
+	dockerPort               = 2376
 )
 
 type Driver struct {
@@ -283,7 +284,7 @@ func (d *Driver) GetURL() (string, error) {
 	if ip == "" {
 		return "", nil
 	}
-	return fmt.Sprintf("tcp://%s:2376", ip), nil
+	return fmt.Sprintf("tcp://%s:%d", ip, dockerPort), nil
 }
 
 func (d *Driver) GetIP() (string, error) {
@@ -528,28 +529,57 @@ func (d *Driver) configureSecurityGroup() error {
 
 	d.SecurityGroupId = securityGroup.GroupId
 
-	perms := []amz.IpPermission{
-		{
+	log.Debugf("configuring authorization %s", ipRange)
+
+	perms := configureSecurityGroupPermissions(securityGroup)
+
+	if len(perms) != 0 {
+		if err := d.getClient().AuthorizeSecurityGroup(d.SecurityGroupId, perms); err != nil {
+			return err
+		}
+
+	}
+
+	return nil
+}
+
+func configureSecurityGroupPermissions(group *amz.SecurityGroup) []amz.IpPermission {
+	hasSshPort := false
+	hasDockerPort := false
+	for _, p := range group.IpPermissions {
+		switch p.FromPort {
+		case 22:
+			hasSshPort = true
+		case dockerPort:
+			hasDockerPort = true
+		}
+	}
+
+	perms := []amz.IpPermission{}
+
+	if !hasSshPort {
+		perm := amz.IpPermission{
 			Protocol: "tcp",
 			FromPort: 22,
 			ToPort:   22,
 			IpRange:  ipRange,
-		},
-		{
+		}
+
+		perms = append(perms, perm)
+	}
+
+	if !hasDockerPort {
+		perm := amz.IpPermission{
 			Protocol: "tcp",
-			FromPort: 2376,
-			ToPort:   2376,
+			FromPort: dockerPort,
+			ToPort:   dockerPort,
 			IpRange:  ipRange,
-		},
+		}
+
+		perms = append(perms, perm)
 	}
 
-	log.Debugf("authorizing %s", ipRange)
-
-	if err := d.getClient().AuthorizeSecurityGroup(d.SecurityGroupId, perms); err != nil {
-		return err
-	}
-
-	return nil
+	return perms
 }
 
 func (d *Driver) deleteSecurityGroup() error {
diff --git a/drivers/amazonec2/amazonec2_test.go b/drivers/amazonec2/amazonec2_test.go
new file mode 100644
index 0000000000..9ea1161f45
--- /dev/null
+++ b/drivers/amazonec2/amazonec2_test.go
@@ -0,0 +1,94 @@
+package amazonec2
+
+import (
+	"testing"
+
+	"github.com/docker/machine/drivers/amazonec2/amz"
+)
+
+var (
+	securityGroup = amz.SecurityGroup{
+		GroupName: "test-group",
+		GroupId:   "12345",
+		VpcId:     "12345",
+	}
+)
+
+const (
+	testSshPort    = 22
+	testDockerPort = 2376
+)
+
+func TestConfigureSecurityGroupPermissionsEmpty(t *testing.T) {
+	group := securityGroup
+	perms := configureSecurityGroupPermissions(&group)
+	if len(perms) != 2 {
+		t.Fatalf("expected 2 permissions; received %d", len(perms))
+	}
+}
+
+func TestConfigureSecurityGroupPermissionsSshOnly(t *testing.T) {
+	group := securityGroup
+
+	group.IpPermissions = []amz.IpPermission{
+		{
+			Protocol: "tcp",
+			FromPort: testSshPort,
+			ToPort:   testSshPort,
+		},
+	}
+
+	perms := configureSecurityGroupPermissions(&group)
+	if len(perms) != 1 {
+		t.Fatalf("expected 1 permission; received %d", len(perms))
+	}
+
+	receivedPort := perms[0].FromPort
+	if receivedPort != testDockerPort {
+		t.Fatalf("expected permission on port %d; received port %d", testDockerPort, receivedPort)
+	}
+}
+
+func TestConfigureSecurityGroupPermissionsDockerOnly(t *testing.T) {
+	group := securityGroup
+
+	group.IpPermissions = []amz.IpPermission{
+		{
+			Protocol: "tcp",
+			FromPort: testDockerPort,
+			ToPort:   testDockerPort,
+		},
+	}
+
+	perms := configureSecurityGroupPermissions(&group)
+	if len(perms) != 1 {
+		t.Fatalf("expected 1 permission; received %d", len(perms))
+	}
+
+	receivedPort := perms[0].FromPort
+	if receivedPort != testSshPort {
+		t.Fatalf("expected permission on port %d; received port %d", testSshPort, receivedPort)
+	}
+}
+
+func TestConfigureSecurityGroupPermissionsDockerAndSsh(t *testing.T) {
+	group := securityGroup
+
+	group.IpPermissions = []amz.IpPermission{
+		{
+			Protocol: "tcp",
+			FromPort: testSshPort,
+			ToPort:   testSshPort,
+		},
+		{
+			Protocol: "tcp",
+			FromPort: testDockerPort,
+			ToPort:   testDockerPort,
+		},
+	}
+
+	perms := configureSecurityGroupPermissions(&group)
+	if len(perms) != 0 {
+		t.Fatalf("expected 0 permissions; received %d", len(perms))
+	}
+}
diff --git a/drivers/amazonec2/amz/security_group.go b/drivers/amazonec2/amz/security_group.go
index 0593643241..c01d3605ed 100644
--- a/drivers/amazonec2/amz/security_group.go
+++ b/drivers/amazonec2/amz/security_group.go
@@ -12,7 +12,9 @@ type DeleteSecurityGroupResponse struct {
 }
 
 type SecurityGroup struct {
-	GroupName string `xml:"groupName"`
-	GroupId   string `xml:"groupId"`
-	VpcId     string `xml:"vpcId"`
+	GroupName           string         `xml:"groupName"`
+	GroupId             string         `xml:"groupId"`
+	VpcId               string         `xml:"vpcId"`
+	IpPermissions       []IpPermission `xml:"ipPermissions,omitempty"`
+	IpPermissionsEgress []IpPermission `xml:"ipPermissionsEgress,omitempty"`
 }