Add support for merging split attributes

This adds BaseNL80211Attr::Merge() and NL80211Packet::
GetAllAttributes(). These functions are useful for merging
attributes across different netlink packets.

This also adds the corresponding unit tests and fixes a existing
typo.

Bug: 36899490
Test: compile, unit tests
Change-Id: Ie53b0e088aee7aab80cf9b06b2f0e822c0062a3d

6 files changed, 142 insertions, 5 deletions

diff --git a/net/nl80211_attribute.cpp b/net/nl80211_attribute.cpp
index c87a879..e61386a 100644
--- a/net/nl80211_attribute.cpp
+++ b/net/nl80211_attribute.cpp
@@ -31,6 +31,14 @@ template class NL80211Attr<vector<uint8_t>>;
 template class NL80211Attr<string>;
 
 // For BaseNL80211Attr
+
+BaseNL80211Attr::BaseNL80211Attr(int id,
+    const vector<uint8_t>& raw_buffer) {
+  size_t size = raw_buffer.size();
+  InitHeaderAndResize(id, size);
+  memcpy(data_.data() + NLA_HDRLEN, raw_buffer.data(), raw_buffer.size());
+}
+
 void BaseNL80211Attr::InitHeaderAndResize(int attribute_id,
                                           int payload_length) {
   data_.resize(NLA_HDRLEN + NLA_ALIGN(payload_length), 0);
@@ -83,12 +91,41 @@ bool BaseNL80211Attr::GetAttributeImpl(const uint8_t* buf,
 }
 
 
+bool BaseNL80211Attr::Merge(const BaseNL80211Attr& other_attr) {
+  if (!other_attr.IsValid()) {
+    LOG(ERROR) << "Can not merge invalid attribute";
+    return false;
+  }
+  if (GetAttributeId() != other_attr.GetAttributeId()) {
+    LOG(ERROR) << "Can not merge attributes with different ids";
+    return false;
+  }
+
+  auto our_header = reinterpret_cast<nlattr*>(data_.data());
+  int our_len_without_padding = our_header->nla_len;
+  auto other_header =
+      reinterpret_cast<const nlattr*>(other_attr.GetConstData().data());
+  int other_len_without_padding = other_header->nla_len;
+  // Update the length to including the content of |attr|.
+  our_header->nla_len =
+      our_len_without_padding + other_len_without_padding - NLA_HDRLEN;
+
+  // Remove padding 0s.
+  data_.resize(our_len_without_padding);
+  // Insert content of |attr|.
+  data_.insert(
+      data_.end(),
+      reinterpret_cast<const uint8_t*>(other_header) + NLA_HDRLEN,
+      reinterpret_cast<const uint8_t*>(other_header) +
+          other_len_without_padding);
+  // Add padding 0s.
+  data_.resize(NLA_ALIGN(our_header->nla_len), 0);
+  return true;
+}
+
 // For NL80211Attr<std::vector<uint8_t>>
 NL80211Attr<vector<uint8_t>>::NL80211Attr(int id,
-    const vector<uint8_t>& raw_buffer) {
-  size_t size = raw_buffer.size();
-  InitHeaderAndResize(id, size);
-  memcpy(data_.data() + NLA_HDRLEN, raw_buffer.data(), raw_buffer.size());
+    const vector<uint8_t>& raw_buffer) : BaseNL80211Attr(id, raw_buffer) {
 }
 
 NL80211Attr<vector<uint8_t>>::NL80211Attr(
diff --git a/net/nl80211_attribute.h b/net/nl80211_attribute.h
index ebae05a..9b0fcc7 100644
--- a/net/nl80211_attribute.h
+++ b/net/nl80211_attribute.h
@@ -32,6 +32,7 @@ namespace wificond {
 
 class BaseNL80211Attr {
  public:
+  BaseNL80211Attr(int id, const std::vector<uint8_t>& raw_buffer);
   virtual ~BaseNL80211Attr() = default;
 
   const std::vector<uint8_t>& GetConstData() const;
@@ -48,6 +49,10 @@ class BaseNL80211Attr {
                               int attr_id,
                               uint8_t** attr_start,
                               uint8_t** attr_end);
+  // Merge the payload of |attr| to current attribute.
+  // This is only used for merging attribute from the response of split dump.
+  // Returns true on success.
+  bool Merge(const BaseNL80211Attr& attr);
 
  protected:
   BaseNL80211Attr() = default;
diff --git a/net/nl80211_packet.cpp b/net/nl80211_packet.cpp
index 0100621..900aadd 100644
--- a/net/nl80211_packet.cpp
+++ b/net/nl80211_packet.cpp
@@ -18,6 +18,8 @@
 
 #include <android-base/logging.h>
 
+using std::make_unique;
+using std::unique_ptr;
 using std::vector;
 
 namespace android {
@@ -215,6 +217,25 @@ bool NL80211Packet::GetAttribute(int id,
   return true;
 }
 
+bool NL80211Packet::GetAllAttributes(
+    vector<BaseNL80211Attr>* attributes) const {
+  const uint8_t* ptr = data_.data() + NLMSG_HDRLEN + GENL_HDRLEN;
+  const uint8_t* end_ptr = data_.data() + data_.size();
+  while (ptr + NLA_HDRLEN <= end_ptr) {
+    auto header = reinterpret_cast<const nlattr*>(ptr);
+    if (ptr + NLA_ALIGN(header->nla_len) > end_ptr ||
+      header->nla_len == 0) {
+      LOG(ERROR) << "broken nl80211 atrribute.";
+      return false;
+    }
+    attributes->emplace_back(
+        header->nla_type,
+        vector<uint8_t>(ptr + NLA_HDRLEN, ptr + header->nla_len));
+    ptr += NLA_ALIGN(header->nla_len);
+  }
+  return true;
+}
+
 void NL80211Packet::DebugLog() const {
   const uint8_t* ptr = data_.data() + NLMSG_HDRLEN + GENL_HDRLEN;
   const uint8_t* end_ptr = data_.data() + data_.size();
diff --git a/net/nl80211_packet.h b/net/nl80211_packet.h
index ed9f02e..4448373 100644
--- a/net/nl80211_packet.h
+++ b/net/nl80211_packet.h
@@ -112,6 +112,11 @@ class NL80211Packet {
 
   bool HasAttribute(int id) const;
   bool GetAttribute(int id, NL80211NestedAttr* attribute) const;
+  // Get all attributes to |*attribute| as a vector.
+  // In case of failure, attributes up until the first invalid attribute
+  // actually will be present in |attributes|.
+  bool GetAllAttributes(
+      std::vector<BaseNL80211Attr>* attributes) const;
 
   template <typename T>
   bool GetAttributeValue(int id, T* value) const {
diff --git a/tests/nl80211_attribute_unittest.cpp b/tests/nl80211_attribute_unittest.cpp
index 24ec92e..c7e1331 100644
--- a/tests/nl80211_attribute_unittest.cpp
+++ b/tests/nl80211_attribute_unittest.cpp
@@ -282,5 +282,44 @@ TEST(NL80211AttributeTest, GetListOfNestedAttributesFromBuffer) {
   EXPECT_TRUE(value3 == 8);
 }
 
+TEST(NL80211AttributeTest, MergeAttributes) {
+  NL80211Attr<std::vector<uint8_t>> attr1(1, {'a', 'b', 'c'});
+  NL80211Attr<std::vector<uint8_t>> attr2(1, {'d', 'e'});
+  ASSERT_TRUE(attr1.Merge(attr2));
+  std::vector<uint8_t> expected_value{{'a', 'b', 'c', 'd', 'e'}};
+  EXPECT_EQ(expected_value, attr1.GetValue());
+}
+
+TEST(NL80211AttributeTest, CannotMergeInvalidAttributeWithBrokenBuffer) {
+  NL80211Attr<std::vector<uint8_t>> valid_attr(1, {'a', 'b', 'c'});
+  std::vector<uint8_t> broken_buffer(
+      kBrokenBuffer,
+      kBrokenBuffer + sizeof(kBrokenBuffer));
+  NL80211Attr<std::vector<uint8_t>> invalid_attr(broken_buffer);
+  EXPECT_FALSE(valid_attr.Merge(invalid_attr));
+}
+
+TEST(NL80211AttributeTest, CannotMergeAttributesWithDifferentIds) {
+  NL80211Attr<std::vector<uint8_t>> attr1(1, {'a', 'b', 'c'});
+  NL80211Attr<std::vector<uint8_t>> attr2(2, {'d', 'e', 'f'});
+  EXPECT_FALSE(attr1.Merge(attr2));
+}
+
+TEST(NL80211AttributeTest, MergeNestedAttributes) {
+  NL80211NestedAttr nested_attr1(0);
+  NL80211NestedAttr nested_attr2(0);
+  NL80211Attr<uint32_t> uint32_attr1(1, kU32Value1);
+  NL80211Attr<uint32_t> uint32_attr2(2, kU32Value2);
+  nested_attr1.AddAttribute(uint32_attr1);
+  nested_attr2.AddAttribute(uint32_attr2);
+  ASSERT_TRUE(nested_attr1.Merge(nested_attr2));
+
+  uint32_t value1, value2;
+  EXPECT_TRUE(nested_attr1.GetAttributeValue(1, &value1));
+  EXPECT_TRUE(value1 == kU32Value1);
+  EXPECT_TRUE(nested_attr1.GetAttributeValue(2, &value2));
+  EXPECT_TRUE(value2 == kU32Value2);
+}
+
 }  // namespace wificond
 }  // namespace android
diff --git a/tests/nl80211_packet_unittest.cpp b/tests/nl80211_packet_unittest.cpp
index b1db4ac..a5fad1b 100644
--- a/tests/nl80211_packet_unittest.cpp
+++ b/tests/nl80211_packet_unittest.cpp
@@ -171,7 +171,7 @@ TEST(NL80211PacketTest, AddNestedAttributesToNL80211Packet) {
   EXPECT_FALSE(netlink_packet.HasAttribute(3));
 }
 
-TEST(NL80211PacketTest, CanbotGetMissingAttributeFromNL80211Packet) {
+TEST(NL80211PacketTest, CannotGetMissingAttributeFromNL80211Packet) {
   NL80211Packet netlink_packet(kNLMsgType,
                                kGenNLCommand,
                                kNLMsgSequenceNumber,
@@ -184,6 +184,36 @@ TEST(NL80211PacketTest, CanbotGetMissingAttributeFromNL80211Packet) {
   EXPECT_FALSE(netlink_packet.GetAttributeValue(2, &attr_value));
 }
 
+TEST(NL80211PacketTest, CanGetAllOfAttributeFromNL80211Packet) {
+  NL80211Packet netlink_packet(kNLMsgType,
+                               kGenNLCommand,
+                               kNLMsgSequenceNumber,
+                               kPortId);
+  NL80211Attr<uint8_t> u8_attr(1, kU8Value1);
+  NL80211Attr<uint32_t> u32_attr_1(2, kU32Value1);
+  NL80211Attr<uint32_t> u32_attr_2(4, kU32Value2);
+  netlink_packet.AddAttribute(u8_attr);
+  netlink_packet.AddAttribute(u32_attr_1);
+  netlink_packet.AddAttribute(u32_attr_2);
+  EXPECT_TRUE(netlink_packet.IsValid());
+  std::vector<BaseNL80211Attr> attributes;
+  EXPECT_TRUE(netlink_packet.GetAllAttributes(&attributes));
+
+  EXPECT_TRUE(attributes.size() == 3);
+
+  NL80211Attr<uint8_t>* u8_attr_retrieved =
+      static_cast<NL80211Attr<uint8_t>*>(&attributes[0]);
+  EXPECT_TRUE(u8_attr_retrieved->GetValue() == kU8Value1);
+
+  NL80211Attr<uint32_t>* u32_attr_1_retrieved =
+      static_cast<NL80211Attr<uint32_t>*>(&attributes[1]);
+  EXPECT_TRUE(u32_attr_1_retrieved->GetValue() == kU32Value1);
+
+  NL80211Attr<uint32_t>* u32_attr_2_retrieved =
+      static_cast<NL80211Attr<uint32_t>*>(&attributes[2]);
+  EXPECT_TRUE(u32_attr_2_retrieved->GetValue() == kU32Value2);
+}
+
 TEST(NL80211PacketTest, ParseCMDAssociateTest) {
   NL80211Packet netlink_packet(std::vector<uint8_t>(
       kNL80211_CMD_ASSOCIATE,