Move ip rule functionality to Python.

This allows us to experiment with ip rules without having to
compile a modified ip binary for the VM architecture.

Change-Id: Ibb0bc510fc2b1686712b94fffae6f1cea26a82bf

3 files changed, 140 insertions, 4 deletions

diff --git a/tests/net_test/iproute.py b/tests/net_test/iproute.py
new file mode 100644
index 00000000..c05c3ff0
--- /dev/null
+++ b/tests/net_test/iproute.py
@@ -0,0 +1,132 @@
+#!/usr/bin/python
+
+"""Partial Python implementation of iproute functionality."""
+
+import os
+import socket
+import struct
+
+### Base netlink constants. See include/uapi/linux/netlink.h.
+NETLINK_ROUTE = 0
+
+# Request constants.
+NLM_F_REQUEST = 1
+NLM_F_ACK = 4
+NLM_F_EXCL = 0x200
+NLM_F_CREATE = 0x400
+
+# Message types.
+NLMSG_ERROR = 2
+
+# Data structure formats.
+STRUCT_NLMSGHDR = "=LHHLL"
+STRUCT_NLMSGERR = "=i"
+STRUCT_NLATTR = "=HH"
+
+
+### rtnetlink constants. See include/uapi/linux/rtnetlink.h.
+# Message types.
+RTM_NEWRULE = 32
+RTM_DELRULE = 33
+
+# Routing message type values (rtm_type).
+RTN_UNSPEC = 0
+RTN_UNICAST = 1
+
+# Routing protocol values (rtm_protocol).
+RTPROT_STATIC = 4
+
+# Route scope values (rtm_scope).
+RT_SCOPE_UNIVERSE = 0
+
+# Data structure formats.
+STRUCT_RTMSG = "=BBBBBBBBI"
+
+
+### FIB rule constants. See include/uapi/linux/fib_rules.h.
+FRA_FWMARK = 10
+FRA_TABLE = 15
+
+
+def Unpack(fmt, data):
+  """Unpacks a data structure with variable-size contents at the end."""
+  size = struct.calcsize(fmt)
+  data, remainder = data[:size], data[size:]
+  return struct.unpack(fmt, data), remainder
+
+
+class IPRoute(object):
+
+  """Provides a tiny subset of iproute functionality."""
+
+  BUFSIZE = 1024
+
+  def _NlAttrU32(self, nla_type, value):
+    data = struct.pack("=I", value)
+    nla_len = struct.calcsize(STRUCT_NLATTR) + len(data)
+    return struct.pack(STRUCT_NLATTR, nla_len, nla_type) + data
+
+  def __init__(self):
+    # Global sequence number.
+    self.seq = 0
+    self.sock = socket.socket(socket.AF_NETLINK, socket.SOCK_RAW,
+                              socket.NETLINK_ROUTE)
+    self.sock.connect((0, 0))  # The kernel.
+    self.pid = self.sock.getsockname()[1]
+
+  def _Send(self, msg):
+    self.seq += 1
+    self.sock.send(msg)
+
+  def _Recv(self):
+    return self.sock.recv(self.BUFSIZE)
+
+  def _Rule(self, version, is_add, table, match_nlattr):
+    """Python equivalent of "ip rule <add|del> <match_cond> lookup <table>".
+
+    Args:
+      version: An integer, 4 or 6.
+      is_add: True to add a rule, False to delete it.
+      table: The table to add/delete the rule from.
+      match_nlattr: A blob of struct nlattrs that express the match condition.
+
+    Raises:
+      IOError: If the netlink request returns an error.
+      ValueError: If the kernel's response could not be parsed.
+    """
+    self.seq += 1
+
+    # Create a struct rtmsg specifying the table and the given match attributes.
+    family = {4: socket.AF_INET, 6: socket.AF_INET6}[version]
+    rtmsg = struct.pack(STRUCT_RTMSG, family, 0, 0, 0, 0,
+                        RTPROT_STATIC, RT_SCOPE_UNIVERSE, RTN_UNICAST, 0)
+    rtmsg += match_nlattr
+    rtmsg += self._NlAttrU32(FRA_TABLE, table)
+
+    # Create a netlink request containing the rtmsg.
+    command = RTM_NEWRULE if is_add else RTM_DELRULE
+    flags = NLM_F_REQUEST | NLM_F_ACK
+    if is_add:
+      flags |= (NLM_F_EXCL | NLM_F_CREATE)
+
+    # Fill in the length field.
+    length = struct.calcsize(STRUCT_NLMSGHDR) + len(rtmsg)
+    nlmsg = struct.pack(STRUCT_NLMSGHDR, length, command, flags,
+                        self.seq, self.pid) + rtmsg
+
+    # Send the message and block forever until we receive a response.
+    self._Send(nlmsg)
+    response = self._Recv()
+
+    # Find the error code.
+    (_, msgtype, _, _, _), msg = Unpack(STRUCT_NLMSGHDR, response)
+    if msgtype == NLMSG_ERROR:
+      ((error,), _) = Unpack(STRUCT_NLMSGERR, msg)
+      if error:
+        raise IOError(error, os.strerror(-error))
+    else:
+      raise ValueError("Unexpected netlink ACK type %d" % msgtype)
+
+  def FwmarkRule(self, version, is_add, fwmark, table):
+    nlattr = self._NlAttrU32(FRA_FWMARK, fwmark)
+    return self._Rule(version, is_add, table, nlattr)
diff --git a/tests/net_test/mark_test.py b/tests/net_test/mark_test.py
index 3ab4b8fb..f6223dcc 100755
--- a/tests/net_test/mark_test.py
+++ b/tests/net_test/mark_test.py
@@ -10,8 +10,10 @@ from socket import *  # pylint: disable=wildcard-import
 import struct
 import time
 import unittest
+
 from scapy import all as scapy
 
+import iproute
 import net_test
 
 DEBUG = False
@@ -226,7 +228,6 @@ class MarkTest(net_test.NetworkTest):
   COMMANDS = [
       "/sbin/%(iptables)s %(append_delete)s INPUT -t mangle -i %(iface)s"
       " -j MARK --set-mark %(netid)d",
-      "ip -%(version)d rule %(add_del)s fwmark %(netid)s lookup %(table)s",
   ]
   ROUTE_COMMANDS = [
       "ip -%(version)d route %(add_del)s table %(table)s"
@@ -243,6 +244,9 @@ class MarkTest(net_test.NetworkTest):
     iface = cls._GetInterfaceName(netid)
     for version, iptables in zip([4, 6], ["iptables", "ip6tables"]):
 
+      table = cls._TableForNetid(netid)
+      cls.iproute.FwmarkRule(version, is_add, netid, table)
+
       if version == 6:
         cmds = cls.COMMANDS
         if cls.AUTOCONF_TABLE_OFFSET < 0:
@@ -267,7 +271,7 @@ class MarkTest(net_test.NetworkTest):
           "macaddr": cls._RouterMacAddress(netid),
           "netid": netid,
           "router": cls._RouterAddress(netid, version),
-          "table": cls._TableForNetid(netid),
+          "table": table,
           "version": version,
       }).split("\n")
       for cmd in cmds:
@@ -310,6 +314,7 @@ class MarkTest(net_test.NetworkTest):
     # This is per-class setup instead of per-testcase setup because shelling out
     # to ip and iptables is slow, and because routing configuration doesn't
     # change during the test.
+    cls.iproute = iproute.IPRoute()
     cls.tuns = {}
     cls.ifindices = {}
     cls._SetAutoconfTableSysctl(1000)
@@ -653,6 +658,5 @@ class MarkTest(net_test.NetworkTest):
       open("/proc/sys/net/ipv4/tcp_syncookies", "w").write("1")
 
 
-
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/net_test/net_test.py b/tests/net_test/net_test.py
index 5017c16c..3686f3e1 100755
--- a/tests/net_test/net_test.py
+++ b/tests/net_test/net_test.py
@@ -208,7 +208,7 @@ def GetDefaultRoute(version=6):
 
 
 def GetDefaultRouteInterface():
-  _, iface = GetDefaultRoute()
+  unused_gw, iface = GetDefaultRoute()
   return iface