Merge pull request #55 from eko/rework-network

Rework network to allow IP generation range from 127.1.2.1 to 127.255.255.255

Author: eko

pkg/proxy/network.go

@@ -24,18 +24,17 @@ func init() {
 func getNetworkInterface() (string, error) {
 	ifaces, err := net.Interfaces()
 	if err != nil {
-		fmt.Println("Cannot retrieve interfaces list: ", err)
-		return "", err
+		return "", fmt.Errorf("cannot retrieve interfaces list: %v", err)
 	}
 
-	for _, i := range ifaces {
-		ifaceFlags := i.Flags.String()
+	for _, iface := range ifaces {
+		ifaceFlags := iface.Flags.String()
 		if strings.Contains(ifaceFlags, "loopback") {
-			return i.Name, nil
+			return iface.Name, nil
 		}
 	}
 
-	return "", errors.New("Unable to find loopback network interface")
+	return "", errors.New("unable to find 'loopback' network interface")
 }
 
 // getAddIPCommandWithArgs returns the command (ifconfig, ip, ...) that will be used for the current OS
@@ -67,52 +66,64 @@ func getAddIPCommandWithArgs(ip string) (string, []string) {
 	return command, args
 }
 
-func generateIP(a byte, b byte, c byte, d int, port string) (net.IP, error) {
-	ip := net.IPv4(a, b, c, byte(d))
-
+func assignIpToPort(a, b, c, d byte, port string) (byte, byte, byte, byte, error) {
 	// Retrieve network interface
 	iface, err := net.InterfaceByName(networkInterface)
 	if err != nil {
-		return net.IP{}, err
+		return a, b, c, d, err
 	}
 
-	command, args := getAddIPCommandWithArgs(ip.String())
+	for {
+		// Maximum IP bytes reached
+		if b == 255 && c == 255 && d == 255 {
+			break
+		}
 
-	for i := d; i < 255; i++ {
-		ip = net.IPv4(a, b, c, byte(i))
+		ip := net.IPv4(a, b, c, d)
 
 		addrs, err := iface.Addrs()
 		if err != nil {
-			return net.IP{}, err
+			return a, b, c, d, err
 		}
 
-		// Try to assign port to the IP addresses already assigned to the interface
-		for _, addr := range addrs {
-			if addr.String() == ip.String()+"/8" {
-				conn, err := net.Dial("tcp", fmt.Sprintf("%s:%s", ip.String(), port))
-				if err != nil {
-					return net.IPv4(a, b, c, byte(i)), nil
-				}
-				if conn != nil {
-					conn.Close()
-				}
+		// In case IP is already assigned to network interface, don't try to create it again
+		if !isAlreadyAssigned(ip, addrs) {
+			command, args := getAddIPCommandWithArgs(ip.String())
+
+			if err := exec.Command(command, args...).Run(); err != nil {
+				return a, b, c, d, fmt.Errorf("error while trying to run ifconfig/ip command to add new IP address (%s) on network interface '%s': %v", ip.String(), networkInterface, err)
 			}
 		}
 
-		// No already assigned IP/Port available, add IP address to the network interface
-		if err := exec.Command(command, args...).Run(); err != nil {
-			return net.IP{}, fmt.Errorf("Cannot run ifconfig command to add new IP address (%s) on network interface '%s': %v", ip.String(), networkInterface, err)
+		// Can't be contacted on ip/port? it means this couple is free to be used
+		if !canDial(ip.String(), port) {
+			return a, b, c, d, nil
 		}
 
-		// Try to assign port to the newly assigned IP
-		conn, err := net.Dial("tcp", fmt.Sprintf("%s:%s", ip.String(), port))
-		if err != nil {
-			return net.IPv4(a, b, c, byte(i)), nil
-		}
-		if conn != nil {
-			conn.Close()
+		a, b, c, d = getNextIPAddress(a, b, c, d)
+	}
+
+	return a, b, c, d, fmt.Errorf("unable to find an available IP/Port (ip: %d.%d.%d.%d:%s)", a, b, c, d, port)
+}
+
+func isAlreadyAssigned(ip net.IP, addrs []net.Addr) bool {
+	for _, addr := range addrs {
+		if addr.String() == ip.String()+"/8" {
+			return true
 		}
 	}
 
-	return net.IP{}, fmt.Errorf("Unable to find an available IP/Port (ip: %d.%d.%d.%d:%s)", a, b, c, d, port)
+	return false
+}
+
+func canDial(ip, port string) bool {
+	conn, err := net.Dial("tcp", fmt.Sprintf("%s:%s", ip, port))
+	if conn != nil {
+		conn.Close()
+	}
+	if err != nil {
+		return false
+	}
+
+	return true
 }

pkg/proxy/proxy.go

@@ -32,8 +32,10 @@ type proxy struct {
 	addProxyForwardMux sync.Mutex
 	listenerMux        sync.Mutex
 	latestPort         string
-	attributedIPs      map[string]net.IP
-	ipLastByte         int
+	lastIpByteA        byte
+	lastIpByteB        byte
+	lastIpByteC        byte
+	lastIpByteD        byte
 	view               ui.View
 }
 
@@ -45,8 +47,10 @@ func NewProxy(view ui.View, hostfile hostfile.Hostfile) *proxy {
 		listeners:     make(map[string]net.Listener),
 		listening:     true,
 		latestPort:    ProxyPortStart,
-		attributedIPs: make(map[string]net.IP, 0),
-		ipLastByte:    1,
+		lastIpByteA:   127,
+		lastIpByteB:   0,
+		lastIpByteC:   1,
+		lastIpByteD:   0,
 		view:          view,
 	}
 }
@@ -160,27 +164,43 @@ func (p *proxy) AddProxyForward(name string, proxyForward *ProxyForward) {
 }
 
 func (p *proxy) generateIP(pf *ProxyForward) error {
-	// Create new listener on a dedicated IP address if the service
-	// does not already have an IP address attributed, elsewhere give the already
-	// attributed address
-	var localIP net.IP
 	var err error
-	if v, ok := p.attributedIPs[pf.Name]; ok {
-		localIP = v
-	} else {
-		localIP, err = generateIP(127, 1, 2, p.ipLastByte, pf.LocalPort)
-		p.ipLastByte = p.ipLastByte + 1
-		if err != nil {
-			return err
-		}
+
+	a, b, c, d := getNextIPAddress(p.lastIpByteA, p.lastIpByteB, p.lastIpByteC, p.lastIpByteD)
+
+	p.lastIpByteA = a
+	p.lastIpByteB = b
+	p.lastIpByteC = c
+	p.lastIpByteD = d
+
+	a, b, c, d, err = assignIpToPort(a, b, c, d, pf.LocalPort)
+	if err != nil {
+		return err
 	}
 
-	p.attributedIPs[pf.Name] = localIP
-	pf.SetLocalIP(localIP.String())
+	ip := net.IPv4(a, b, c, d)
+	pf.SetLocalIP(ip.String())
 
 	return nil
 }
 
+func getNextIPAddress(a, b, c, d byte) (byte, byte, byte, byte) {
+	if b == 255 && c == 255 && d == 255 {
+		return a, b, c, d
+	} else if c == 255 && d == 255 {
+		b++
+		c = 1
+		d = 1
+	} else if d == 255 {
+		c++
+		d = 1
+	} else {
+		d++
+	}
+
+	return a, b, c, d
+}
+
 func (p *proxy) generateProxyPort(proxyForward *ProxyForward) {
 	integerPort, _ := strconv.Atoi(p.latestPort)
 	p.latestPort = strconv.Itoa(integerPort + 1)

pkg/proxy/proxy_test.go

@@ -3,6 +3,7 @@
 package proxy
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/eko/monday/pkg/hostfile"
@@ -29,7 +30,7 @@ func TestNewProxy(t *testing.T) {
 
 	assert.Len(t, p.ProxyForwards, 0)
 	assert.Equal(t, p.latestPort, "9400")
-	assert.Equal(t, p.ipLastByte, 1)
+	assert.Equal(t, p.ipLastByte, byte(1))
 }
 
 func TestAddProxyForward(t *testing.T) {
@@ -53,7 +54,7 @@ func TestAddProxyForward(t *testing.T) {
 	// Then
 	assert.Len(t, proxy.ProxyForwards, 1)
 	assert.Equal(t, proxy.latestPort, "9401")
-	assert.Equal(t, proxy.ipLastByte, 2)
+	assert.Equal(t, proxy.ipLastByte, byte(2))
 }
 
 func TestAddProxyForwardWhenMultiple(t *testing.T) {
@@ -120,5 +121,75 @@ func TestListen(t *testing.T) {
 
 	assert.Len(t, proxy.ProxyForwards, 1)
 	assert.Equal(t, proxy.latestPort, "9401")
-	assert.Equal(t, proxy.ipLastByte, 2)
+	assert.Equal(t, proxy.ipLastByte, byte(2))
+}
+
+func TestGetNextIPAddress(t *testing.T) {
+	testCases := []struct {
+		a byte
+		b byte
+		c byte
+		d byte
+
+		expectedA byte
+		expectedB byte
+		expectedC byte
+		expectedD byte
+	}{
+		{ // Case incrementing d
+			a: 127, b: 0, c: 0, d: 1,
+			expectedA: 127, expectedB: 0, expectedC: 0, expectedD: 2,
+		},
+		{ // Case incrementing d to last byte
+			a: 127, b: 0, c: 0, d: 254,
+			expectedA: 127, expectedB: 0, expectedC: 0, expectedD: 255,
+		},
+		{ // Case incrementing c and reset d to 1
+			a: 127, b: 0, c: 0, d: 255,
+			expectedA: 127, expectedB: 0, expectedC: 1, expectedD: 1,
+		},
+		{ // Case incrementing d to last byte
+			a: 127, b: 0, c: 1, d: 254,
+			expectedA: 127, expectedB: 0, expectedC: 1, expectedD: 255,
+		},
+		{ // Case incrementing c and reset d to 1
+			a: 127, b: 0, c: 1, d: 255,
+			expectedA: 127, expectedB: 0, expectedC: 2, expectedD: 1,
+		},
+		{ // Case incrementing d to last byte
+			a: 127, b: 0, c: 254, d: 254,
+			expectedA: 127, expectedB: 0, expectedC: 254, expectedD: 255,
+		},
+		{ // Case incrementing c and reset d to 1
+			a: 127, b: 0, c: 254, d: 255,
+			expectedA: 127, expectedB: 0, expectedC: 255, expectedD: 1,
+		},
+		{ // Case incrementing d to last byte when c is already on latest byte
+			a: 127, b: 0, c: 255, d: 254,
+			expectedA: 127, expectedB: 0, expectedC: 255, expectedD: 255,
+		},
+		{ // Case incrementing b and reset c and d to last byte
+			a: 127, b: 0, c: 255, d: 255,
+			expectedA: 127, expectedB: 1, expectedC: 1, expectedD: 1,
+		},
+		{ // Case incrementing d to last byte when b and c are already on latest byte
+			a: 127, b: 255, c: 255, d: 254,
+			expectedA: 127, expectedB: 255, expectedC: 255, expectedD: 255,
+		},
+		{ // Reached max level
+			a: 127, b: 255, c: 255, d: 255,
+			expectedA: 127, expectedB: 255, expectedC: 255, expectedD: 255,
+		},
+	}
+
+	for i, testCase := range testCases {
+		t.Run(fmt.Sprintf("Case %d", i), func(t *testing.T) {
+			a, b, c, d := getNextIPAddress(testCase.a, testCase.b, testCase.c, testCase.d)
+
+			assert.Equal(t, testCase.expectedA, a)
+			assert.Equal(t, testCase.expectedB, b)
+			assert.Equal(t, testCase.expectedC, c)
+			assert.Equal(t, testCase.expectedD, d)
+		})
+	}
 }