From ac2a4e64106b28d7b05d1ee31716af365424ff83 Mon Sep 17 00:00:00 2001
From: Michael Crosby <michael@crosbymichael.com>
Date: Thu, 23 Jan 2014 07:46:42 -0800
Subject: [PATCH] Refactor to support multiple ip addresses

Docker-DCO-1.1-Signed-off-by: Michael Crosby <michael@crosbymichael.com> (github: crosbymichael)
---
 network.go                                   |  3 +-
 networkdriver/portallocator/portallocator.go | 91 ++++++++++++++++----
 2 files changed, 73 insertions(+), 21 deletions(-)

diff --git a/network.go b/network.go
index 3ec1f6fb73..ddb0af5bc7 100644
--- a/network.go
+++ b/network.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"github.com/dotcloud/docker/networkdriver"
 	"github.com/dotcloud/docker/networkdriver/ipallocator"
+	"github.com/dotcloud/docker/networkdriver/portallocator"
 	"github.com/dotcloud/docker/pkg/iptables"
 	"github.com/dotcloud/docker/pkg/netlink"
 	"github.com/dotcloud/docker/proxy"
@@ -20,8 +21,6 @@ const (
 	DefaultNetworkBridge = "docker0"
 	DisableNetworkBridge = "none"
 	DefaultNetworkMtu    = 1500
-	portRangeStart       = 49153
-	portRangeEnd         = 65535
 	siocBRADDBR          = 0x89a0
 )
 
diff --git a/networkdriver/portallocator/portallocator.go b/networkdriver/portallocator/portallocator.go
index a0a3ebd8d2..5c7ad04caa 100644
--- a/networkdriver/portallocator/portallocator.go
+++ b/networkdriver/portallocator/portallocator.go
@@ -3,11 +3,19 @@ package portallocator
 import (
 	"errors"
 	"github.com/dotcloud/docker/pkg/collections"
+	"net"
 	"sync"
 )
 
 type portMappings map[string]*collections.OrderedIntSet
 
+type ipData struct {
+	allocatedPorts portMappings
+	availablePorts portMappings
+}
+
+type ipMapping map[net.IP]*ipData
+
 const (
 	BeginPortRange = 49153
 	EndPortRange   = 65535
@@ -20,22 +28,62 @@ var (
 )
 
 var (
-	lock           = sync.Mutex{}
-	allocatedPorts = portMappings{}
-	availablePorts = portMappings{}
+	defaultIPData *ipData
+
+	lock      = sync.Mutex{}
+	ips       = ipMapping{}
+	defaultIP = net.ParseIP("0.0.0.0")
 )
 
 func init() {
-	allocatedPorts["udp"] = collections.NewOrderedIntSet()
-	availablePorts["udp"] = collections.NewOrderedIntSet()
-	allocatedPorts["tcp"] = collections.NewOrderedIntSet()
-	availablePorts["tcp"] = collections.NewOrderedIntSet()
+	defaultIPData = newIpData()
+	ips[defaultIP] = defaultIP
+}
+
+func newIpData() {
+	data := &ipData{
+		allocatedPorts: portMappings{},
+		availablePorts: portMappings{},
+	}
+
+	data.allocatedPorts["udp"] = collections.NewOrderedIntSet()
+	data.availablePorts["udp"] = collections.NewOrderedIntSet()
+	data.allocatedPorts["tcp"] = collections.NewOrderedIntSet()
+	data.availablePorts["tcp"] = collections.NewOrderedIntSet()
+
+	return data
+}
+
+func getData(ip net.IP) *ipData {
+	data, exists := ips[ip]
+	if !exists {
+		data = newIpData()
+		ips[ip] = data
+	}
+	return data
+}
+
+func validateMapping(data *ipData, proto string, port int) error {
+	allocated := data.allocatedPorts[proto]
+	if allocated.Exists(proto) {
+		return ErrPortAlreadyAllocated
+	}
+	return nil
+}
+
+func usePort(data *ipData, proto string, port int) {
+	allocated, available := data.allocatedPorts[proto], data.availablePorts[proto]
+	for i := 0; i < 2; i++ {
+		allocated.Push(port)
+		available.Remove(port)
+		allocated, available = defaultIPData.allocatedPorts[proto], defaultIPData.availablePorts[proto]
+	}
 }
 
 // RequestPort returns an available port if the port is 0
 // If the provided port is not 0 then it will be checked if
 // it is available for allocation
-func RequestPort(proto string, port int) (int, error) {
+func RequestPort(ip net.IP, proto string, port int) (int, error) {
 	lock.Lock()
 	defer lock.Unlock()
 
@@ -43,20 +91,28 @@ func RequestPort(proto string, port int) (int, error) {
 		return 0, err
 	}
 
-	var (
-		allocated = allocatedPorts[proto]
-		available = availablePorts[proto]
-	)
+	data := getData(ip)
+	allocated, available := data.allocatedPorts[proto], data.availablePorts[proto]
 
+	// If the user requested a specific port to be allocated
 	if port != 0 {
-		if allocated.Exists(port) {
-			return 0, ErrPortAlreadyAllocated
+		if err := validateMapping(defaultIP, proto, port); err != nil {
+			return 0, err
 		}
+
+		if !defaultIP.Equal(ip) {
+			if err := validateMapping(data, proto, port); err != nil {
+				return 0, err
+			}
+		}
+
 		available.Remove(port)
 		allocated.Push(port)
+
 		return port, nil
 	}
 
+	// Dynamic allocation
 	next := available.Pop()
 	if next == 0 {
 		next = allocated.PullBack()
@@ -76,7 +132,7 @@ func RequestPort(proto string, port int) (int, error) {
 
 // ReleasePort will return the provided port back into the
 // pool for reuse
-func ReleasePort(proto string, port int) error {
+func ReleasePort(ip net.IP, proto string, port int) error {
 	lock.Lock()
 	defer lock.Unlock()
 
@@ -84,10 +140,7 @@ func ReleasePort(proto string, port int) error {
 		return err
 	}
 
-	var (
-		allocated = allocatedPorts[proto]
-		available = availablePorts[proto]
-	)
+	allocated, available := getCollection(ip, proto)
 
 	allocated.Remove(port)
 	available.Push(port)