Merge changes I6919795b,I91c5ddfb,I7c3f4d73

am: 68d2588030

Change-Id: Ib1c081af6398c161ba8aecb8cbd68164711bf7be

5 files changed, 111 insertions, 36 deletions

diff --git a/net/test/cstruct.py b/net/test/cstruct.py
index 91cd72e..43c47a2 100644
--- a/net/test/cstruct.py
+++ b/net/test/cstruct.py
@@ -17,7 +17,10 @@
 Example usage:
 
 >>> # Declare a struct type by specifying name, field formats and field names.
-... # Field formats are the same as those used in the struct module.
+... # Field formats are the same as those used in the struct module, except:
+... # - S: Nested Struct.
+... # - A: NULL-padded ASCII string. Like s, but printing ignores contiguous
+... #      trailing NULL blocks at the end.
 ... import cstruct
 >>> NLMsgHdr = cstruct.Struct("NLMsgHdr", "=LHHLL", "length type flags seq pid")
 >>>
@@ -44,6 +47,15 @@ NLMsgHdr(length=44, type=33, flags=2, seq=0, pid=510)
 >>> cstruct.Read(data, NLMsgHdr)
 (NLMsgHdr(length=44, type=33, flags=2, seq=0, pid=510), 'more data')
 >>>
+>>> # Structs can contain one or more nested structs. The nested struct types
+... # are specified in a list as an optional last argument. Nested structs may
+... # contain nested structs.
+... S = cstruct.Struct("S", "=BI", "byte1 int2")
+>>> N = cstruct.Struct("N", "!BSiS", "byte1 s2 int3 s2", [S, S])
+>>> NN = cstruct.Struct("NN", "SHS", "s1 word2 n3", [S, N])
+>>> nn = NN((S((1, 25000)), -29876, N((55, S((5, 6)), 1111, S((7, 8))))))
+>>> nn.n3.s2.int2 = 5
+>>>
 """
 
 import ctypes
@@ -51,10 +63,18 @@ import string
 import struct
 
 
+def CalcSize(fmt):
+  if "A" in fmt:
+    fmt = fmt.replace("A", "s")
+  return struct.calcsize(fmt)
+
 def CalcNumElements(fmt):
-  size = struct.calcsize(fmt)
+  prevlen = len(fmt)
+  fmt = fmt.replace("S", "")
+  numstructs = prevlen - len(fmt)
+  size = CalcSize(fmt)
   elements = struct.unpack(fmt, "\x00" * size)
-  return len(elements)
+  return len(elements) + numstructs
 
 
 def Struct(name, fmt, fieldnames, substructs={}):
@@ -80,6 +100,8 @@ def Struct(name, fmt, fieldnames, substructs={}):
     _fieldnames = fieldnames
     # Dict mapping field indices to nested struct classes.
     _nested = {}
+    # List of string fields that are ASCII strings.
+    _asciiz = set()
 
     if isinstance(_fieldnames, str):
       _fieldnames = _fieldnames.split(" ")
@@ -95,11 +117,16 @@ def Struct(name, fmt, fieldnames, substructs={}):
         _nested[index] = substructs[laststructindex]
         laststructindex += 1
         _format += "%ds" % len(_nested[index])
+      elif fmt[i] == "A":
+        # Null-terminated ASCII string.
+        index = CalcNumElements(fmt[:i])
+        _asciiz.add(index)
+        _format += "s"
       else:
          # Standard struct format character.
         _format += fmt[i]
 
-    _length = struct.calcsize(_format)
+    _length = CalcSize(_format)
 
     def _SetValues(self, values):
       super(CStruct, self).__setattr__("_values", list(values))
@@ -165,9 +192,11 @@ def Struct(name, fmt, fieldnames, substructs={}):
 
     def __str__(self):
       def FieldDesc(index, name, value):
-        if isinstance(value, str) and any(
-            c not in string.printable for c in value):
-          value = value.encode("hex")
+        if isinstance(value, str):
+          if index in self._asciiz:
+            value = value.rstrip("\x00")
+          elif any(c not in string.printable for c in value):
+            value = value.encode("hex")
         return "%s=%s" % (name, value)
 
       descriptions = [
diff --git a/net/test/cstruct_test.py b/net/test/cstruct_test.py
index 2d5a408..fdcbd55 100755
--- a/net/test/cstruct_test.py
+++ b/net/test/cstruct_test.py
@@ -55,6 +55,61 @@ class CstructTest(unittest.TestCase):
       self.CheckEquals(i, i)
     self.CheckEquals(a1, a3)
 
+  def testNestedStructs(self):
+    Nested = cstruct.Struct("Nested", "!HSSi",
+                            "word1 nest2 nest3 int4",
+                            [TestStructA, TestStructB])
+    DoubleNested = cstruct.Struct("DoubleNested", "SSB",
+                                  "nest1 nest2 byte3",
+                                  [TestStructA, Nested])
+    d = DoubleNested((TestStructA((1, 2)),
+                      Nested((5, TestStructA((3, 4)), TestStructB((7, 8)), 9)),
+                      6))
+
+    expectedlen = (len(TestStructA) +
+                   2 + len(TestStructA) + len(TestStructB) + 4 +
+                   1)
+    self.assertEquals(expectedlen, len(DoubleNested))
+
+    self.assertEquals(7, d.nest2.nest3.byte1)
+
+    d.byte3 = 252
+    d.nest2.word1 = 33214
+    n = d.nest2
+    n.int4 = -55
+    t = n.nest3
+    t.int2 = 33627591
+
+    self.assertEquals(33627591, d.nest2.nest3.int2)
+
+    expected = (
+        "DoubleNested(nest1=TestStructA(byte1=1, int2=2),"
+        " nest2=Nested(word1=33214, nest2=TestStructA(byte1=3, int2=4),"
+        " nest3=TestStructB(byte1=7, int2=33627591), int4=-55), byte3=252)")
+    self.assertEquals(expected, str(d))
+    expected = ("01" "02000000"
+                "81be" "03" "04000000"
+                "07" "c71d0102" "ffffffc9" "fc").decode("hex")
+    self.assertEquals(expected, d.Pack())
+    unpacked = DoubleNested(expected)
+    self.CheckEquals(unpacked, d)
+
+  def testNullTerminatedStrings(self):
+    TestStruct = cstruct.Struct("TestStruct", "B16si16AH",
+                                "byte1 string2 int3 ascii4 word5")
+    nullstr = "hello" + (16 - len("hello")) * "\x00"
+
+    t = TestStruct((2, nullstr, 12345, nullstr, 33210))
+    expected = ("TestStruct(byte1=2, string2=68656c6c6f0000000000000000000000,"
+                " int3=12345, ascii4=hello, word5=33210)")
+    self.assertEquals(expected, str(t))
+
+    embeddednull = "hello\x00visible123"
+    t = TestStruct((2, embeddednull, 12345, embeddednull, 33210))
+    expected = ("TestStruct(byte1=2, string2=68656c6c6f0076697369626c65313233,"
+                " int3=12345, ascii4=hello\x00visible123, word5=33210)")
+    self.assertEquals(expected, str(t))
+
 
 if __name__ == "__main__":
   unittest.main()
diff --git a/net/test/iproute.py b/net/test/iproute.py
index 2451a68..1d980af 100644
--- a/net/test/iproute.py
+++ b/net/test/iproute.py
@@ -31,18 +31,6 @@ import netlink
 ### Base netlink constants. See include/uapi/linux/netlink.h.
 NETLINK_ROUTE = 0
 
-# Request constants.
-NLM_F_REQUEST = 1
-NLM_F_ACK = 4
-NLM_F_REPLACE = 0x100
-NLM_F_EXCL = 0x200
-NLM_F_CREATE = 0x400
-NLM_F_DUMP = 0x300
-
-# Message types.
-NLMSG_ERROR = 2
-NLMSG_DONE = 3
-
 # Data structure formats.
 # These aren't constants, they're classes. So, pylint: disable=invalid-name
 NLMsgHdr = cstruct.Struct("NLMsgHdr", "=LHHLL", "length type flags seq pid")
@@ -285,7 +273,7 @@ class IPRoute(netlink.NetlinkSocket):
                   "IFA_LABEL"]:
       data = nla_data.strip("\x00")
     elif name == "RTA_METRICS":
-      data = self._ParseAttributes(-RTA_METRICS, msg.family, None, nla_data)
+      data = self._ParseAttributes(-RTA_METRICS, None, nla_data)
     elif name == "RTA_CACHEINFO":
       data = RTACacheinfo(nla_data)
     elif name == "IFA_CACHEINFO":
@@ -310,12 +298,12 @@ class IPRoute(netlink.NetlinkSocket):
   def _SendNlRequest(self, command, data, flags=0):
     """Sends a netlink request and expects an ack."""
 
-    flags |= NLM_F_REQUEST
+    flags |= netlink.NLM_F_REQUEST
     if CommandVerb(command) != "GET":
-      flags |= NLM_F_ACK
+      flags |= netlink.NLM_F_ACK
     if CommandVerb(command) == "NEW":
-      if not flags & NLM_F_REPLACE:
-        flags |= (NLM_F_EXCL | NLM_F_CREATE)
+      if not flags & netlink.NLM_F_REPLACE:
+        flags |= (netlink.NLM_F_EXCL | netlink.NLM_F_CREATE)
 
     super(IPRoute, self)._SendNlRequest(command, data, flags)
 
@@ -403,7 +391,7 @@ class IPRoute(netlink.NetlinkSocket):
     except IndexError:
       raise ValueError("Don't know how to print command type %s" % name)
 
-  def MaybeDebugCommand(self, command, data):
+  def MaybeDebugCommand(self, command, unused_flags, data):
     subject = CommandSubject(command)
     if "ALL" not in self.NL_DEBUG and subject not in self.NL_DEBUG:
       return
@@ -500,7 +488,7 @@ class IPRoute(netlink.NetlinkSocket):
     self._Route(version, RTM_GETROUTE, 0, dest, prefixlen, None, oif, mark, uid)
     data = self._Recv()
     # The response will either be an error or a list of routes.
-    if NLMsgHdr(data).type == NLMSG_ERROR:
+    if NLMsgHdr(data).type == netlink.NLMSG_ERROR:
       self._ParseAck(data)
     routes = self._GetMsgList(RTMsg, data, False)
     return routes
@@ -531,7 +519,7 @@ class IPRoute(netlink.NetlinkSocket):
 
   def UpdateNeighbour(self, version, addr, lladdr, dev, state):
     self._Neighbour(version, True, addr, lladdr, dev, state,
-                    flags=NLM_F_REPLACE)
+                    flags=netlink.NLM_F_REPLACE)
 
   def DumpNeighbours(self, version):
     ndmsg = NdMsg((self._AddressFamily(version), 0, 0, 0, 0))
diff --git a/net/test/netlink.py b/net/test/netlink.py
index 261350b..ee57b20 100644
--- a/net/test/netlink.py
+++ b/net/test/netlink.py
@@ -91,7 +91,7 @@ class NetlinkSocket(object):
     """No-op, nonspecific version of decode."""
     return nla_type, nla_data
 
-  def _ParseAttributes(self, command, family, msg, data):
+  def _ParseAttributes(self, command, msg, data):
     """Parses and decodes netlink attributes.
 
     Takes a block of NLAttr data structures, decodes them using Decode, and
@@ -99,7 +99,6 @@ class NetlinkSocket(object):
 
     Args:
       command: An integer, the rtnetlink command being carried out.
-      family: The address family.
       msg: A Struct, the type of the data after the netlink header.
       data: A byte string containing a sequence of NLAttr data structures.
 
@@ -139,6 +138,10 @@ class NetlinkSocket(object):
     self.sock.connect((0, 0))  # The kernel.
     self.pid = self.sock.getsockname()[1]
 
+  def MaybeDebugCommand(self, command, flags, data):
+    # Default no-op implementation to be overridden by subclasses.
+    pass
+
   def _Send(self, msg):
     # self._Debug(msg.encode("hex"))
     self.seq += 1
@@ -174,7 +177,7 @@ class NetlinkSocket(object):
     length = len(NLMsgHdr) + len(data)
     nlmsg = NLMsgHdr((length, command, flags, self.seq, self.pid)).Pack()
 
-    self.MaybeDebugCommand(command, nlmsg + data)
+    self.MaybeDebugCommand(command, flags, nlmsg + data)
 
     # Send the message.
     self._Send(nlmsg + data)
@@ -196,8 +199,7 @@ class NetlinkSocket(object):
 
     # Parse the attributes in the nlmsg.
     attrlen = nlmsghdr.length - len(nlmsghdr) - len(nlmsg)
-    attributes = self._ParseAttributes(nlmsghdr.type, nlmsg.family,
-                                       nlmsg, data[:attrlen])
+    attributes = self._ParseAttributes(nlmsghdr.type, nlmsg, data[:attrlen])
     data = data[attrlen:]
     return (nlmsg, attributes), data
 
@@ -223,7 +225,7 @@ class NetlinkSocket(object):
 
     Args:
       command: An integer, the command to run (e.g., RTM_NEWADDR).
-      msg: A string, the raw bytes of the request (e.g., a packed RTMsg).
+      msg: A struct, the request (e.g., a RTMsg). May be None.
       msgtype: A cstruct.Struct, the data type to parse the dump results as.
       attrs: A string, the raw bytes of any request attributes to include.
 
@@ -233,12 +235,13 @@ class NetlinkSocket(object):
     """
     # Create a netlink dump request containing the msg.
     flags = NLM_F_DUMP | NLM_F_REQUEST
+    msg = "" if msg is None else msg.Pack()
     length = len(NLMsgHdr) + len(msg) + len(attrs)
     nlmsghdr = NLMsgHdr((length, command, flags, self.seq, self.pid))
 
     # Send the request.
-    request = nlmsghdr.Pack() + msg.Pack() + attrs
-    self.MaybeDebugCommand(command, request)
+    request = nlmsghdr.Pack() + msg + attrs
+    self.MaybeDebugCommand(command, flags, request)
     self._Send(request)
 
     # Keep reading netlink messages until we get a NLMSG_DONE.
diff --git a/net/test/sock_diag.py b/net/test/sock_diag.py
index 13ed0a8..1865891 100755
--- a/net/test/sock_diag.py
+++ b/net/test/sock_diag.py
@@ -160,7 +160,7 @@ class SockDiag(netlink.NetlinkSocket):
 
     return name, data
 
-  def MaybeDebugCommand(self, command, data):
+  def MaybeDebugCommand(self, command, unused_flags, data):
     name = self._GetConstantName(__name__, command, "SOCK_")
     if "ALL" not in self.NL_DEBUG and "SOCK" not in self.NL_DEBUG:
       return