patchpanel: Query iptables counters in CountersService

This patch adds the code for querying iptables and parsing its output to
get traffic counters. Also adds some corresponding unit tests.

BUG=b:160113164
TEST=cros_workon_make --board=$BOARD --test patchpanel

Change-Id: Ifd73c077c8796dc8f84dfe1301bf57140d30f58a
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform2/+/2291854
Tested-by: Jie Jiang <jiejiang@chromium.org>
Commit-Queue: Jie Jiang <jiejiang@chromium.org>
Reviewed-by: Hugo Benichi <hugobenichi@google.com>
diff --git a/patchpanel/BUILD.gn b/patchpanel/BUILD.gn
index 4ba6de3..3f72524 100644
--- a/patchpanel/BUILD.gn
+++ b/patchpanel/BUILD.gn
@@ -33,6 +33,7 @@
     "libshill-client",
     "libshill-net-${libbase_ver}",
     "protobuf-lite",
+    "re2",
     "system_api",
   ]
   if (use.fuzzer) {
diff --git a/patchpanel/counters_service.cc b/patchpanel/counters_service.cc
index 091b3bd..8a1bfda 100644
--- a/patchpanel/counters_service.cc
+++ b/patchpanel/counters_service.cc
@@ -6,16 +6,116 @@
 
 #include <set>
 #include <string>
+#include <utility>
 #include <vector>
 
+#include <base/strings/strcat.h>
+#include <base/strings/string_split.h>
+#include <re2/re2.h>
+
 namespace patchpanel {
 
 namespace {
 
+using Counter = CountersService::Counter;
+using SourceDevice = CountersService::SourceDevice;
+
 constexpr char kMangleTable[] = "mangle";
 
+// The following regexs and code is written and tested for iptables v1.6.2.
+// Output code of iptables can be found at:
+//   https://git.netfilter.org/iptables/tree/iptables/iptables.c?h=v1.6.2
+
+// The chain line looks like:
+//   "Chain tx_fwd_eth0 (1 references)".
+// This regex extracts "tx" (direction), "eth0" (ifname) from this example, and
+// "fwd" (prebuilt chain) is matched but not captured because it is not required
+// for counters.
+constexpr LazyRE2 kChainLine = {
+    R"(Chain (rx|tx)_(?:input|fwd|postrt)_(\w+).*)"};
+
+// The counter line looks like (some spaces are deleted to make it fit in one
+// line):
+//   "    6511 68041668    all  --  any    any     anywhere   anywhere"
+// The first two counters are captured for pkts and bytes.
+constexpr LazyRE2 kCounterLine = {R"( *(\d+) +(\d+).*)"};
+
+// Parses the output of `iptables -L -x -v` (or `ip6tables`) and adds the parsed
+// values into the corresponding counters in |counters|. An example of |output|
+// can be found in the test file. This function will try to find the pattern of:
+//   <one chain line for an accounting chain>
+//   <one header line>
+//   <one counter line for an accounting rule>
+// The interface name and direction (rx or tx) will be extracted from the chain
+// line, and then the values extracted from the counter line will be added into
+// the counter for that interface. Note that this function will not fully
+// validate if |output| is an output from iptables.
+bool ParseOutput(const std::string& output,
+                 const std::set<std::string>& devices,
+                 std::map<SourceDevice, Counter>* counters) {
+  DCHECK(counters);
+  const std::vector<std::string> lines = base::SplitString(
+      output, "\n", base::KEEP_WHITESPACE, base::SPLIT_WANT_ALL);
+
+  // Finds the chain line for an accounting chain first, and then parse the
+  // following line(s) to get the counters for this chain. Repeats this process
+  // until we reach the end of |output|.
+  for (auto it = lines.cbegin(); it != lines.cend(); it++) {
+    // Finds the chain name line.
+    std::string direction, ifname;
+    while (it != lines.cend() &&
+           !RE2::FullMatch(*it, *kChainLine, &direction, &ifname))
+      it++;
+
+    if (it == lines.cend())
+      break;
+
+    // Skips this group if this ifname is not requested.
+    if (!devices.empty() && devices.find(ifname) == devices.end())
+      continue;
+
+    // Skips the chain name line and the header line.
+    if (lines.cend() - it <= 2) {
+      LOG(ERROR) << "Invalid iptables output";
+      return false;
+    }
+    it += 2;
+
+    // The current line should be the accounting rule containing the counters.
+    // Currently we only have one accounting rule (UNKNOWN source) for each
+    // chain.
+    // TODO(jiejiang): The following part will be extended to a loop when we
+    // have more accounting rules.
+    uint64_t pkts, bytes;
+    if (!RE2::FullMatch(*it, *kCounterLine, &pkts, &bytes)) {
+      LOG(ERROR) << "Cannot parse \"" << *it << "\"";
+      return false;
+    }
+
+    TrafficCounter::Source source = TrafficCounter::UNKNOWN;
+    auto& counter = (*counters)[std::make_pair(source, ifname)];
+    if (direction == "rx") {
+      counter.rx_packets += pkts;
+      counter.rx_bytes += bytes;
+    } else {
+      counter.tx_packets += pkts;
+      counter.tx_bytes += bytes;
+    }
+  }
+  return true;
+}
+
 }  // namespace
 
+Counter::Counter(uint64_t rx_bytes,
+                 uint64_t rx_packets,
+                 uint64_t tx_bytes,
+                 uint64_t tx_packets)
+    : rx_bytes(rx_bytes),
+      rx_packets(rx_packets),
+      tx_bytes(tx_bytes),
+      tx_packets(tx_packets) {}
+
 CountersService::CountersService(ShillClient* shill_client,
                                  MinijailedProcessRunner* runner)
     : shill_client_(shill_client), runner_(runner) {
@@ -25,6 +125,39 @@
       &CountersService::OnDeviceChanged, weak_factory_.GetWeakPtr()));
 }
 
+std::map<SourceDevice, Counter> CountersService::GetCounters(
+    const std::set<std::string>& devices) {
+  std::map<SourceDevice, Counter> counters;
+
+  // Handles counters for IPv4 and IPv6 separately and returns failure if either
+  // of the procession fails, since counters for only IPv4 or IPv6 are biased.
+  std::string iptables_result;
+  int ret = runner_->iptables(kMangleTable, {"-L", "-x", "-v", "-w"},
+                              true /*log_failures*/, &iptables_result);
+  if (ret != 0 || iptables_result.empty()) {
+    LOG(ERROR) << "Failed to query IPv4 counters";
+    return {};
+  }
+  if (!ParseOutput(iptables_result, devices, &counters)) {
+    LOG(ERROR) << "Failed to parse IPv4 counters";
+    return {};
+  }
+
+  std::string ip6tables_result;
+  ret = runner_->ip6tables(kMangleTable, {"-L", "-x", "-v", "-w"},
+                           true /*log_failures*/, &ip6tables_result);
+  if (ret != 0 || ip6tables_result.empty()) {
+    LOG(ERROR) << "Failed to query IPv6 counters";
+    return {};
+  }
+  if (!ParseOutput(ip6tables_result, devices, &counters)) {
+    LOG(ERROR) << "Failed to parse IPv6 counters";
+    return {};
+  }
+
+  return counters;
+}
+
 void CountersService::OnDeviceChanged(const std::set<std::string>& added,
                                       const std::set<std::string>& removed) {
   for (const auto& ifname : added)
diff --git a/patchpanel/counters_service.h b/patchpanel/counters_service.h
index 2876c2e..b2e8cde 100644
--- a/patchpanel/counters_service.h
+++ b/patchpanel/counters_service.h
@@ -5,10 +5,14 @@
 #ifndef PATCHPANEL_COUNTERS_SERVICE_H_
 #define PATCHPANEL_COUNTERS_SERVICE_H_
 
+#include <map>
 #include <set>
 #include <string>
+#include <utility>
 #include <vector>
 
+#include <patchpanel/proto_bindings/patchpanel_service.pb.h>
+
 #include "patchpanel/minijailed_process_runner.h"
 #include "patchpanel/shill_client.h"
 
@@ -43,17 +47,39 @@
 // The above rules and chains will never be removed once created, so we will
 // check if one rule exists before creating it.
 //
-// TODO(jiejiang): Query will be implemented in future patches.
-//
 // Query: Two commands (iptables and ip6tables) will be executed in the mangle
 // table to get all the chains and rules. And then we perform a text parsing on
 // the output to get the counters. Counters for the same entry will be merged
 // before return.
 class CountersService {
  public:
+  using SourceDevice = std::pair<TrafficCounter::Source, std::string>;
+  struct Counter {
+    Counter() = default;
+    Counter(uint64_t rx_bytes,
+            uint64_t rx_packets,
+            uint64_t tx_bytes,
+            uint64_t tx_packets);
+
+    uint64_t rx_bytes = 0;
+    uint64_t rx_packets = 0;
+    uint64_t tx_bytes = 0;
+    uint64_t tx_packets = 0;
+  };
+
   CountersService(ShillClient* shill_client, MinijailedProcessRunner* runner);
   ~CountersService() = default;
 
+  // Collects and returns counters from all the existing iptables rules.
+  // |devices| is the set of interfaces for which counters should be returned,
+  // any unknown interfaces will be ignored. If |devices| is empty, counters for
+  // all known interfaces will be returned. An empty map will be returned on
+  // any failure. Note that currently all traffic to/from an interface will be
+  // counted by (UNKNOWN, ifname), i.e., no other sources except for UNKNOWN are
+  // used.
+  std::map<SourceDevice, Counter> GetCounters(
+      const std::set<std::string>& devices);
+
  private:
   // TODO(b/161060333): Move the following two functions elsewhere.
   // Creates a new chain using both iptables and ip6tables in the mangle table.
diff --git a/patchpanel/counters_service_test.cc b/patchpanel/counters_service_test.cc
index 5b9e77d..2e6f904 100644
--- a/patchpanel/counters_service_test.cc
+++ b/patchpanel/counters_service_test.cc
@@ -15,10 +15,106 @@
 #include "patchpanel/fake_shill_client.h"
 
 namespace patchpanel {
-namespace {
 
+using ::testing::ContainerEq;
 using ::testing::Contains;
+using ::testing::DoAll;
 using ::testing::ElementsAreArray;
+using ::testing::Return;
+using ::testing::SetArgPointee;
+
+using Counter = CountersService::Counter;
+using SourceDevice = CountersService::SourceDevice;
+
+// The following two functions should be put outside the anounymous namespace
+// otherwise they could not be found in the tests.
+std::ostream& operator<<(std::ostream& os, const Counter& counter) {
+  os << "rx_bytes:" << counter.rx_bytes << ", rx_packets:" << counter.rx_packets
+     << ", tx_bytes:" << counter.tx_bytes
+     << ", tx_packets:" << counter.tx_packets;
+  return os;
+}
+
+bool operator==(const CountersService::Counter lhs,
+                const CountersService::Counter rhs) {
+  return lhs.rx_bytes == rhs.rx_bytes && lhs.rx_packets == rhs.rx_packets &&
+         lhs.tx_bytes == rhs.tx_bytes && lhs.tx_packets == rhs.tx_packets;
+}
+
+namespace {
+// The following string is copied from the real output of iptables v1.6.2 by
+// `iptables -t mangle -L -x -v`. This output contains all the accounting
+// chains/rules for eth0 and wlan0.
+// TODO(jiejiang): presubmit checker is complaining about the line length for
+// this (and the other raw strings in this file). Find a way to make it happy.
+const char kIptablesOutput[] = R"(
+Chain PREROUTING (policy ACCEPT 22785 packets, 136093545 bytes)
+    pkts      bytes target     prot opt in     out     source               destination
+      18     2196 MARK       all  --  arcbr0 any     anywhere             anywhere             MARK set 0x1
+       0        0 MARK       all  --  vmtap+ any     anywhere             anywhere             MARK set 0x1
+    6526 68051766 MARK       all  --  arc_eth0 any     anywhere             anywhere             MARK set 0x1
+       9     1104 MARK       all  --  arc_wlan0 any     anywhere             anywhere             MARK set 0x1
+
+Chain INPUT (policy ACCEPT 4421 packets, 2461233 bytes)
+    pkts      bytes target     prot opt in     out     source               destination
+  312491 1767147156 rx_input_eth0  all  --  eth0   any     anywhere             anywhere
+       0        0 rx_input_wlan0  all  --  wlan0  any     anywhere             anywhere
+
+Chain FORWARD (policy ACCEPT 18194 packets, 133612816 bytes)
+    pkts      bytes target     prot opt in     out     source               destination
+    6511 68041668 tx_fwd_eth0  all  --  any    eth0    anywhere             anywhere
+   11683 65571148 rx_fwd_eth0  all  --  eth0   any     anywhere             anywhere
+       0        0 tx_fwd_wlan0  all  --  any    wlan0   anywhere             anywhere
+       0        0 rx_fwd_wlan0  all  --  wlan0  any     anywhere             anywhere
+
+Chain OUTPUT (policy ACCEPT 4574 packets, 2900995 bytes)
+    pkts      bytes target     prot opt in     out     source               destination
+
+Chain POSTROUTING (policy ACCEPT 22811 packets, 136518827 bytes)
+    pkts      bytes target     prot opt in     out     source               destination
+  202160 1807550291 tx_postrt_eth0  all  --  any    eth0    anywhere             anywhere             owner socket exists
+       2       96 tx_postrt_wlan0  all  --  any    wlan0   anywhere             anywhere             owner socket exists
+
+Chain tx_fwd_eth0 (1 references)
+    pkts      bytes target     prot opt in     out     source               destination
+    6511 68041668            all  --  any    any     anywhere             anywhere
+
+Chain tx_fwd_wlan0 (1 references)
+    pkts      bytes target     prot opt in     out     source               destination
+       0        0            all  --  any    any     anywhere             anywhere
+
+Chain tx_postrt_eth0 (1 references)
+    pkts      bytes target     prot opt in     out     source               destination
+  202160 1807550291            all  --  any    any     anywhere             anywhere
+
+Chain tx_postrt_wlan0 (1 references)
+    pkts      bytes target     prot opt in     out     source               destination
+       2       96            all  --  any    any     anywhere             anywhere
+
+Chain rx_fwd_eth0 (1 references)
+    pkts      bytes target     prot opt in     out     source               destination
+   11683 65571148            all  --  any    any     anywhere             anywhere
+
+Chain rx_fwd_wlan0 (1 references)
+    pkts      bytes target     prot opt in     out     source               destination
+       0        0            all  --  any    any     anywhere             anywhere
+
+Chain rx_input_eth0 (1 references)
+    pkts      bytes target     prot opt in     out     source               destination
+  312491 1767147156            all  --  any    any     anywhere             anywhere
+
+Chain rx_input_wlan0 (1 references)
+    pkts      bytes target     prot opt in     out     source               destination
+       0        0            all  --  any    any     anywhere             anywhere
+)";
+
+// The expected counters for the above output. "* 2" because the same string
+// will be returned for both iptables and ip6tables in the tests.
+const Counter kCounter_eth0{(65571148 + 1767147156ULL) * 2 /*rx_bytes*/,
+                            (11683 + 312491) * 2 /*rx_packets*/,
+                            (68041668 + 1807550291ULL) * 2 /*tx_bytes*/,
+                            (6511 + 202160) * 2 /*tx_packets*/};
+const Counter kCounter_wlan0{0, 0, 96 * 2, 2 * 2};
 
 class MockProcessRunner : public MinijailedProcessRunner {
  public:
@@ -49,6 +145,34 @@
         std::make_unique<CountersService>(fake_shill_client_.get(), &runner_);
   }
 
+  // Makes `iptables` returning a bad |output|. Expects an empty map from
+  // GetCounters().
+  void TestBadIptablesOutput(const std::string& output) {
+    EXPECT_CALL(runner_, iptables(_, _, _, _))
+        .WillRepeatedly(DoAll(SetArgPointee<3>(output), Return(0)));
+    EXPECT_CALL(runner_, ip6tables(_, _, _, _))
+        .WillRepeatedly(DoAll(SetArgPointee<3>(kIptablesOutput), Return(0)));
+
+    auto actual = counters_svc_->GetCounters({});
+    std::map<SourceDevice, Counter> expected;
+
+    EXPECT_THAT(actual, ContainerEq(expected));
+  }
+
+  // Makes `ip6tables` returning a bad |output|. Expects an empty map from
+  // GetCounters().
+  void TestBadIp6tablesOutput(const std::string& output) {
+    EXPECT_CALL(runner_, iptables(_, _, _, _))
+        .WillRepeatedly(DoAll(SetArgPointee<3>(kIptablesOutput), Return(0)));
+    EXPECT_CALL(runner_, ip6tables(_, _, _, _))
+        .WillRepeatedly(DoAll(SetArgPointee<3>(output), Return(0)));
+
+    auto actual = counters_svc_->GetCounters({});
+    std::map<SourceDevice, Counter> expected;
+
+    EXPECT_THAT(actual, ContainerEq(expected));
+  }
+
   FakeShillClientHelper shill_helper_;
   MockProcessRunner runner_;
   std::unique_ptr<FakeShillClient> fake_shill_client_;
@@ -109,5 +233,82 @@
                                                   brillo::Any(devices));
 }
 
+TEST_F(CountersServiceTest, QueryTrafficCounters) {
+  EXPECT_CALL(runner_, iptables(_, _, _, _))
+      .WillRepeatedly(DoAll(SetArgPointee<3>(kIptablesOutput), Return(0)));
+  EXPECT_CALL(runner_, ip6tables(_, _, _, _))
+      .WillRepeatedly(DoAll(SetArgPointee<3>(kIptablesOutput), Return(0)));
+
+  auto actual = counters_svc_->GetCounters({});
+
+  std::map<SourceDevice, Counter> expected{
+      {{TrafficCounter::UNKNOWN, "eth0"}, kCounter_eth0},
+      {{TrafficCounter::UNKNOWN, "wlan0"}, kCounter_wlan0},
+  };
+
+  EXPECT_THAT(actual, ContainerEq(expected));
+}
+
+TEST_F(CountersServiceTest, QueryTrafficCountersWithFilter) {
+  EXPECT_CALL(runner_, iptables(_, _, _, _))
+      .WillRepeatedly(DoAll(SetArgPointee<3>(kIptablesOutput), Return(0)));
+  EXPECT_CALL(runner_, ip6tables(_, _, _, _))
+      .WillRepeatedly(DoAll(SetArgPointee<3>(kIptablesOutput), Return(0)));
+
+  // Only counters for eth0 should be returned. eth1 should be ignored.
+  auto actual = counters_svc_->GetCounters({"eth0", "eth1"});
+
+  std::map<SourceDevice, Counter> expected{
+      {{TrafficCounter::UNKNOWN, "eth0"}, kCounter_eth0},
+  };
+
+  EXPECT_THAT(actual, ContainerEq(expected));
+}
+
+TEST_F(CountersServiceTest, QueryTrafficCountersWithEmptyIPv4Output) {
+  const std::string kEmptyOutput = "";
+  TestBadIptablesOutput(kEmptyOutput);
+}
+
+TEST_F(CountersServiceTest, QueryTrafficCountersWithEmptyIPv6Output) {
+  const std::string kEmptyOutput = "";
+  TestBadIp6tablesOutput(kEmptyOutput);
+}
+
+TEST_F(CountersServiceTest, QueryTrafficCountersWithOnlyChainName) {
+  const std::string kBadOutput = R"(
+Chain tx_fwd_eth0 (1 references)
+    pkts      bytes target     prot opt in     out     source               destination
+    6511 68041668            all  --  any    any     anywhere             anywhere
+
+Chain tx_fwd_wlan0 (1 references)
+)";
+  TestBadIptablesOutput(kBadOutput);
+}
+
+TEST_F(CountersServiceTest, QueryTrafficCountersWithOnlyChainNameAndHeader) {
+  const std::string kBadOutput = R"(
+Chain tx_fwd_eth0 (1 references)
+    pkts      bytes target     prot opt in     out     source               destination
+    6511 68041668            all  --  any    any     anywhere             anywhere
+
+Chain tx_fwd_wlan0 (1 references)
+    pkts      bytes target     prot opt in     out     source               destination
+)";
+  TestBadIptablesOutput(kBadOutput);
+}
+
+TEST_F(CountersServiceTest, QueryTrafficCountersWithNotFinishedCountersLine) {
+  const std::string kBadOutput = R"(
+Chain tx_fwd_eth0 (1 references)
+    pkts      bytes target     prot opt in     out     source               destination
+    6511 68041668            all  --  any    any     anywhere             anywhere
+
+Chain tx_fwd_wlan0 (1 references)
+    pkts      bytes target     prot opt in     out     source               destination    pkts      bytes target     prot opt in     out     source               destination
+       0     )";
+  TestBadIptablesOutput(kBadOutput);
+}
+
 }  // namespace
 }  // namespace patchpanel