Merge pull request #258 from nak3/ipv4inc

Increment IP addrress safely
aws · Dec 15, 2018 · 87052b1 · 87052b1
2 parents 3020119 + 96a86f5
commit 87052b1
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 11 deletions.
diff --git a/pkg/networkutils/network.go b/pkg/networkutils/network.go
@@ -14,6 +14,7 @@
 package networkutils
 
 import (
+	"encoding/binary"
 	"fmt"
 	"io"
 	"math"
@@ -507,22 +508,24 @@ func setupENINetwork(eniIP string, eniMAC string, eniTable int, eniSubnetCIDR st
 
 	deviceNumber := link.Attrs().Index
 
-	_, gw, err := net.ParseCIDR(eniSubnetCIDR)
+	_, ipnet, err := net.ParseCIDR(eniSubnetCIDR)
 
 	if err != nil {
 		return errors.Wrapf(err, "eni network setup: invalid ipv4 cidr block %s", eniSubnetCIDR)
 	}
 
-	// TODO: big/little endian:  convert subnet to gw
-	gw.IP[3] = gw.IP[3] + 1
+	gw, err := incrementIPv4Addr(ipnet.IP)
+	if err != nil {
+		return errors.Wrapf(err, "eni network setup: failed to define gateway address from %v", ipnet.IP)
+	}
 
-	log.Debugf("Setting up ENI's default gateway %v", gw.IP)
+	log.Debugf("Setting up ENI's default gateway %v", gw)
 
 	for _, r := range []netlink.Route{
 		// Add a direct link route for the host's ENI IP only
 		{
 			LinkIndex: deviceNumber,
-			Dst:       &net.IPNet{IP: gw.IP, Mask: net.CIDRMask(32, 32)},
+			Dst:       &net.IPNet{IP: gw, Mask: net.CIDRMask(32, 32)},
 			Scope:     netlink.SCOPE_LINK,
 			Table:     eniTable,
 		},
@@ -531,7 +534,7 @@ func setupENINetwork(eniIP string, eniMAC string, eniTable int, eniSubnetCIDR st
 			LinkIndex: deviceNumber,
 			Dst:       &net.IPNet{IP: net.IPv4zero, Mask: net.CIDRMask(0, 32)},
 			Scope:     netlink.SCOPE_UNIVERSE,
-			Gw:        gw.IP,
+			Gw:        gw,
 			Table:     eniTable,
 		},
 	} {
@@ -550,13 +553,13 @@ func setupENINetwork(eniIP string, eniMAC string, eniTable int, eniSubnetCIDR st
 					retry++
 					if retry > maxRetryRouteAdd {
 						log.Errorf("Failed to add route %s/0 via %s table %d",
-							r.Dst.IP.String(), gw.IP.String(), eniTable)
+							r.Dst.IP.String(), gw.String(), eniTable)
 						return errors.Wrapf(err,
 							"eni network setup: failed unable to add route %s/0 via %s table %d",
-							r.Dst.IP.String(), gw.IP.String(), eniTable)
+							r.Dst.IP.String(), gw.String(), eniTable)
 					}
 					log.Debugf("Not able to add route route %s/0 via %s table %d (attempt %d/%d)",
-						r.Dst.IP.String(), gw.IP.String(), eniTable,
+						r.Dst.IP.String(), gw.String(), eniTable,
 						retry, maxRetryRouteAdd)
 					time.Sleep(retryRouteAddInterval)
 				} else if isRouteExistsError(err) {
@@ -567,10 +570,10 @@ func setupENINetwork(eniIP string, eniMAC string, eniTable int, eniSubnetCIDR st
 					break
 				} else {
 					return errors.Wrapf(err, "eni network setup: unable to add route %s/0 via %s table %d",
-						r.Dst.IP.String(), gw.IP.String(), eniTable)
+						r.Dst.IP.String(), gw.String(), eniTable)
 				}
 			} else {
-				log.Debugf("Successfully added route route %s/0 via %s table %d", r.Dst.IP.String(), gw.IP.String(), eniTable)
+				log.Debugf("Successfully added route route %s/0 via %s table %d", r.Dst.IP.String(), gw.String(), eniTable)
 				break
 			}
 		}
@@ -598,6 +601,22 @@ func setupENINetwork(eniIP string, eniMAC string, eniTable int, eniSubnetCIDR st
 	return nil
 }
 
+// incremetnIPv4Addr returns incremented IPv4 address
+func incrementIPv4Addr(ip net.IP) (net.IP, error) {
+	ip4 := ip.To4()
+	if ip4 == nil {
+		return nil, fmt.Errorf("%q is not a valid IPv4 Address.", ip)
+	}
+	int_ip := binary.BigEndian.Uint32([]byte(ip4))
+	if int_ip == (1<<32 - 1) {
+		return nil, fmt.Errorf("%q will be overflowed", ip)
+	}
+	int_ip++
+	bytes := make([]byte, 4)
+	binary.BigEndian.PutUint32(bytes, int_ip)
+	return net.IP(bytes), nil
+}
+
 // isNotExistsError returns true if the error type is syscall.ESRCH
 // This helps us determine if we should ignore this error as the route
 // that we want to cleanup has been deleted already routing table

diff --git a/pkg/networkutils/network_test.go b/pkg/networkutils/network_test.go
@@ -196,6 +196,32 @@ func TestSetupHostNetworkNodePortEnabled(t *testing.T) {
 	assert.Equal(t, mockFile{closed: true, data: "2"}, mockRPFilter)
 }
 
+func TestIncrementIPv4Addr(t *testing.T) {
+	testCases := []struct {
+		name     string
+		ip       net.IP
+		expected net.IP
+		err      bool
+	}{
+		{"increment", net.IPv4(10, 0, 0, 1), net.IPv4(10, 0, 0, 2).To4(), false},
+		{"carry up 1", net.IPv4(10, 0, 0, 255), net.IPv4(10, 0, 1, 0).To4(), false},
+		{"carry up 2", net.IPv4(10, 0, 255, 255), net.IPv4(10, 1, 0, 0).To4(), false},
+		{"overflow", net.IPv4(255, 255, 255, 255), nil, true},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			result, err := incrementIPv4Addr(tc.ip)
+			if tc.err {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+			}
+			assert.Equal(t, tc.expected, result, tc.name)
+		})
+	}
+}
+
 type mockIptables struct {
 	// dataplaneState is a map from table name to chain name to slice of rulespecs
 	dataplaneState map[string]map[string][][]string