Add support for IP aliases. (#345)

4 files changed, 53 insertions, 15 deletions

diff --git a/google_compute_engine/ip_forwarding/ip_forwarding_daemon.py b/google_compute_engine/ip_forwarding/ip_forwarding_daemon.py
index a2e43c9..45b0ce7 100755
--- a/google_compute_engine/ip_forwarding/ip_forwarding_daemon.py
+++ b/google_compute_engine/ip_forwarding/ip_forwarding_daemon.py
@@ -84,7 +84,7 @@ class IpForwardingDaemon(object):
     if not to_add and not to_remove:
       return
     self.logger.info(
-        'Changing %s forwarded IPs from %s to %s by adding %s and removing %s.',
+        'Changing %s IPs from %s to %s by adding %s and removing %s.',
         interface, configured or None, desired or None, to_add or None,
         to_remove or None)
 
@@ -128,14 +128,16 @@ class IpForwardingDaemon(object):
     """Called when network interface metadata changes.
 
     Args:
-      result: string, the metadata response with the new network interfaces.
+      result: dict, the metadata response with the new network interfaces.
     """
     for network_interface in result:
       mac_address = network_interface.get('mac')
       interface = self.network_utils.GetNetworkInterface(mac_address)
+      ip_addresses = []
       if interface:
-        forwarded_ips = network_interface.get('forwardedIps')
-        self._HandleForwardedIps(forwarded_ips, interface)
+        ip_addresses.extend(network_interface.get('forwardedIps', []))
+        ip_addresses.extend(network_interface.get('ipAliases', []))
+        self._HandleForwardedIps(ip_addresses, interface)
       else:
         message = 'Network interface not found for MAC address: %s.'
         self.logger.warning(message, mac_address)
diff --git a/google_compute_engine/ip_forwarding/ip_forwarding_utils.py b/google_compute_engine/ip_forwarding/ip_forwarding_utils.py
index c7f6cdf..201d326 100644
--- a/google_compute_engine/ip_forwarding/ip_forwarding_utils.py
+++ b/google_compute_engine/ip_forwarding/ip_forwarding_utils.py
@@ -19,6 +19,7 @@ import re
 import subprocess
 
 IP_REGEX = re.compile(r'\A(\d{1,3}\.){3}\d{1,3}\Z')
+IP_ALIAS_REGEX = re.compile(r'\A(\d{1,3}\.){3}\d{1,3}/\d{1,2}\Z')
 
 
 class IpForwardingUtils(object):
@@ -92,7 +93,7 @@ class IpForwardingUtils(object):
     addresses = []
     forwarded_ips = forwarded_ips or []
     for ip in forwarded_ips:
-      if ip and IP_REGEX.match(ip):
+      if ip and (IP_REGEX.match(ip) or IP_ALIAS_REGEX.match(ip)):
         addresses.append(ip)
       else:
         self.logger.warning('Could not parse IP address: "%s".', ip)
@@ -119,7 +120,8 @@ class IpForwardingUtils(object):
       address: string, the IP address to configure.
       interface: string, the output device to use.
     """
-    args = ['add', 'to', 'local', '%s/32' % address]
+    address = address if IP_ALIAS_REGEX.match(address) else '%s/32' % address
+    args = ['add', 'to', 'local', address]
     options = self._CreateRouteOptions(dev=interface)
     self._RunIpRoute(args=args, options=options)
 
@@ -130,6 +132,7 @@ class IpForwardingUtils(object):
       address: string, the IP address to configure.
       interface: string, the output device to use.
     """
-    args = ['delete', 'to', 'local', '%s/32' % address]
+    address = address if IP_ALIAS_REGEX.match(address) else '%s/32' % address
+    args = ['delete', 'to', 'local', address]
     options = self._CreateRouteOptions(dev=interface)
     self._RunIpRoute(args=args, options=options)
diff --git a/google_compute_engine/ip_forwarding/tests/ip_forwarding_daemon_test.py b/google_compute_engine/ip_forwarding/tests/ip_forwarding_daemon_test.py
index a34e9e8..4c0f32f 100644
--- a/google_compute_engine/ip_forwarding/tests/ip_forwarding_daemon_test.py
+++ b/google_compute_engine/ip_forwarding/tests/ip_forwarding_daemon_test.py
@@ -180,23 +180,26 @@ class IpForwardingDaemonTest(unittest.TestCase):
     mocks.attach_mock(self.mock_network_utils, 'network')
     mocks.attach_mock(self.mock_setup, 'setup')
     self.mock_network_utils.GetNetworkInterface.side_effect = [
-        'eth0', 'eth1', 'eth2', None]
+        'eth0', 'eth1', 'eth2', 'eth3', None]
     result = [
-        {'mac': '1', 'forwardedIps': 'a'},
-        {'mac': '2', 'forwardedIps': 'b'},
-        {'mac': '3'},
-        {'forwardedIps': 'c'},
+        {'mac': '1', 'forwardedIps': ['a']},
+        {'mac': '2', 'forwardedIps': ['b'], 'ipAliases': ['banana']},
+        {'mac': '3', 'ipAliases': ['cherry']},
+        {'mac': '4'},
+        {'forwardedIps': ['d'], 'ipAliases': ['date']},
     ]
 
     ip_forwarding_daemon.IpForwardingDaemon.HandleNetworkInterfaces(
         self.mock_setup, result)
     expected_calls = [
         mock.call.network.GetNetworkInterface('1'),
-        mock.call.setup._HandleForwardedIps('a', 'eth0'),
+        mock.call.setup._HandleForwardedIps(['a'], 'eth0'),
         mock.call.network.GetNetworkInterface('2'),
-        mock.call.setup._HandleForwardedIps('b', 'eth1'),
+        mock.call.setup._HandleForwardedIps(['b', 'banana'], 'eth1'),
         mock.call.network.GetNetworkInterface('3'),
-        mock.call.setup._HandleForwardedIps(None, 'eth2'),
+        mock.call.setup._HandleForwardedIps(['cherry'], 'eth2'),
+        mock.call.network.GetNetworkInterface('4'),
+        mock.call.setup._HandleForwardedIps([], 'eth3'),
         mock.call.network.GetNetworkInterface(None),
         mock.call.setup.logger.warning(mock.ANY, None),
     ]
diff --git a/google_compute_engine/ip_forwarding/tests/ip_forwarding_utils_test.py b/google_compute_engine/ip_forwarding/tests/ip_forwarding_utils_test.py
index 18774a6..5543d13 100644
--- a/google_compute_engine/ip_forwarding/tests/ip_forwarding_utils_test.py
+++ b/google_compute_engine/ip_forwarding/tests/ip_forwarding_utils_test.py
@@ -137,6 +137,12 @@ class IpForwardingUtilsTest(unittest.TestCase):
         '1.1.1.a': False,
         None: False,
         '1.0.0.0': True,
+        '1.1.1.1/1': True,
+        '1.1.1.1/11': True,
+        '123.123.123.123/1': True,
+        '123.123.123.123/123': False,
+        '123.123.123.123/a': False,
+        '123.123.123.123/': False,
     }
     input_ips = forwarded_ips.keys()
     valid_ips = [ip for ip, valid in forwarded_ips.items() if valid]
@@ -175,6 +181,18 @@ class IpForwardingUtilsTest(unittest.TestCase):
     mock_run.assert_called_once_with(
         args=['add', 'to', 'local', '1.1.1.1/32'], options=self.options)
 
+  def testAddIpAlias(self):
+    mock_options = mock.Mock()
+    mock_options.return_value = self.options
+    mock_run = mock.Mock()
+    self.mock_utils._CreateRouteOptions = mock_options
+    self.mock_utils._RunIpRoute = mock_run
+
+    self.mock_utils.AddForwardedIp('1.1.1.1/24', 'interface')
+    mock_options.assert_called_once_with(dev='interface')
+    mock_run.assert_called_once_with(
+        args=['add', 'to', 'local', '1.1.1.1/24'], options=self.options)
+
   def testRemoveForwardedIp(self):
     mock_options = mock.Mock()
     mock_options.return_value = self.options
@@ -187,6 +205,18 @@ class IpForwardingUtilsTest(unittest.TestCase):
     mock_run.assert_called_once_with(
         args=['delete', 'to', 'local', '1.1.1.1/32'], options=self.options)
 
+  def testRemoveAliasIp(self):
+    mock_options = mock.Mock()
+    mock_options.return_value = self.options
+    mock_run = mock.Mock()
+    self.mock_utils._CreateRouteOptions = mock_options
+    self.mock_utils._RunIpRoute = mock_run
+
+    self.mock_utils.RemoveForwardedIp('1.1.1.1/24', 'interface')
+    mock_options.assert_called_once_with(dev='interface')
+    mock_run.assert_called_once_with(
+        args=['delete', 'to', 'local', '1.1.1.1/24'], options=self.options)
+
 
 if __name__ == '__main__':
   unittest.main()