Merge 'goog/master' into gce-dev

Test: TreeHugger
Change-Id: I461a47f02f7dd7b443aa3d1edb0ce2254fef5a11
diff --git a/Android.bp b/Android.bp
index 348a98b..e63080d 100644
--- a/Android.bp
+++ b/Android.bp
@@ -42,12 +42,12 @@
         "cuttlefish_kernel_headers",
     ],
     target: {
-        linux_glibc: {
+        linux: {
             host_ldlibs: ["-lrt"],
             cflags: ["-DCUTTLEFISH_HOST"],
         },
         // We don't need 32 bit host-side builds
-        linux_glibc_x86: {
+        linux_x86: {
             enabled: false,
         },
         // We don't need Darwin host-side builds
@@ -105,7 +105,7 @@
         "liblog",
     ],
     target: {
-        linux_glibc: {
+        linux: {
             srcs: [
                 "host/vsoc/lib/gralloc_buffer_region_view.cpp",
                 "host/vsoc/lib/host_lock.cpp",
@@ -147,7 +147,7 @@
         "libbase",
     ],
     target : {
-        linux_glibc_x86_64: {
+        linux_x86_64: {
             static_libs: [
                 "libcuttlefish_host_config",
                 "libgflags",
diff --git a/common/commands/wificlient/Android.bp b/common/commands/wificlient/Android.bp
index 067d085..cc2e268 100644
--- a/common/commands/wificlient/Android.bp
+++ b/common/commands/wificlient/Android.bp
@@ -5,7 +5,9 @@
     ],
     shared_libs: [
         "libbase",
-        "libcuttlefish_wifi",
+        "vsoc_lib",
+        "libcuttlefish_fs",
+        "cuttlefish_auto_resources",
         "liblog",
         "libnl",
     ],
@@ -15,5 +17,18 @@
     header_libs: [
         "cuttlefish_glog",
     ],
-    defaults: ["cuttlefish_host_and_guest", "cuttlefish_native_isa"],
+    target: {
+        linux: {
+            static_libs: [
+                "libcuttlefish_wifi",
+                "libcuttlefish_host_config",
+            ],
+        },
+        android: {
+            static_libs: [
+                "libcuttlefish_wifi",
+            ],
+        }
+    },
+    defaults: ["cuttlefish_host_and_guest", "cuttlefish_native_isa"]
 }
diff --git a/common/commands/wificlient/main.cc b/common/commands/wificlient/main.cc
index e060408..1b6736b 100644
--- a/common/commands/wificlient/main.cc
+++ b/common/commands/wificlient/main.cc
@@ -18,6 +18,7 @@
 #include <glog/logging.h>
 
 #include "common/libs/wifi/netlink.h"
+#include "common/libs/wifi/packet_switch.h"
 #include "common/libs/wifi/virtual_wifi.h"
 
 DEFINE_string(router, "cvd-wifirouter", "Path to WIFI Router Unix socket.");
@@ -35,6 +36,12 @@
     exit(1);
   }
 
+  cvd::PacketSwitch pktswitch(nl.get());
+  if (!pktswitch.Init()) {
+    LOG(ERROR) << "Could not initialize packet switch.";
+    exit(1);
+  }
+
   std::unique_ptr<cvd::VirtualWIFI> radio(
       new cvd::VirtualWIFI(nl.get(), FLAGS_iface, FLAGS_macaddr));
   if (!radio->Init()) {
@@ -42,5 +49,7 @@
     exit(1);
   }
 
+  pktswitch.Start();
+
   pause();
 }
diff --git a/common/libs/wifi/Android.bp b/common/libs/wifi/Android.bp
index 0ce06f1..6a6e846 100644
--- a/common/libs/wifi/Android.bp
+++ b/common/libs/wifi/Android.bp
@@ -1,10 +1,11 @@
-cc_library_shared {
+cc_library_static {
     name: "libcuttlefish_wifi",
 
     srcs: [
+        "cmd.cc",
         "netlink.cc",
         "nl_client.cc",
-        "cmd.cc",
+        "packet_switch.cc",
         "virtual_wifi.cc",
         "wr_client.cc",
     ],
@@ -16,5 +17,13 @@
     header_libs: [
         "cuttlefish_glog",
     ],
-    defaults: ["cuttlefish_host_and_guest", "cuttlefish_native_isa"],
+    target: {
+        linux: {
+            static_libs: [
+                "libcuttlefish_host_config",
+                "libgflags",
+            ],
+        }
+    },
+    defaults: ["cuttlefish_host_and_guest", "cuttlefish_native_isa"]
 }
diff --git a/common/libs/wifi/cmd.cc b/common/libs/wifi/cmd.cc
index 1a153cc..2576359 100644
--- a/common/libs/wifi/cmd.cc
+++ b/common/libs/wifi/cmd.cc
@@ -19,6 +19,8 @@
 
 Cmd::Cmd() : msg_(nlmsg_alloc()) {}
 
+Cmd::Cmd(nlmsghdr* h) : msg_(nlmsg_convert(h)) {}
+
 Cmd::~Cmd() {
   for (auto& msg : responses_) {
     nlmsg_free(msg);
diff --git a/common/libs/wifi/cmd.h b/common/libs/wifi/cmd.h
index f5d9da0..f9ff509 100644
--- a/common/libs/wifi/cmd.h
+++ b/common/libs/wifi/cmd.h
@@ -28,6 +28,7 @@
 class Cmd {
  public:
   Cmd();
+  explicit Cmd(nlmsghdr* h);
   ~Cmd();
 
   // Cmd() creates netlink request to be sent to kernel.
diff --git a/common/libs/wifi/packet_switch.cc b/common/libs/wifi/packet_switch.cc
new file mode 100644
index 0000000..f365c29
--- /dev/null
+++ b/common/libs/wifi/packet_switch.cc
@@ -0,0 +1,110 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/libs/wifi/packet_switch.h"
+#include "common/libs/wifi/router.h"
+
+#ifdef CUTTLEFISH_HOST
+#include "host/libs/config/host_config.h"
+#endif
+
+namespace cvd {
+
+PacketSwitch::~PacketSwitch() { Stop(); }
+
+bool PacketSwitch::Init() {
+#ifdef CUTTLEFISH_HOST
+  return shm_wifi_.Open(vsoc::GetDomain().c_str());
+#else
+  return shm_wifi_.Open();
+#endif
+}
+
+void PacketSwitch::Start() {
+  std::lock_guard<std::mutex> l(op_mutex_);
+  if (started_) return;
+  // set started to true immediately; this attribute is referenced by threads to
+  // know whether they should terminate.
+  started_ = true;
+
+  nl_->WRCL().SetDefaultHandler(
+      [this](nl_msg* m) { ProcessPacket(m, false); });
+
+  shm_xchg_.reset(new std::thread([this] {
+      size_t maxlen = getpagesize();
+      std::unique_ptr<uint8_t[]> msg(new uint8_t[maxlen]);
+      auto hdr = reinterpret_cast<nlmsghdr*>(msg.get());
+      std::unique_ptr<nl_msg, void (*)(nl_msg*)> nlm(nullptr, nlmsg_free);
+
+      int item = 0;
+      while (started_) {
+        auto len = shm_wifi_.Recv(msg.get(), maxlen);
+        nlm.reset(nlmsg_convert(hdr));
+        ProcessPacket(nlm.get(), true);
+      }
+    }));
+}
+
+void PacketSwitch::Stop() {
+  std::lock_guard<std::mutex> l(op_mutex_);
+  if (!started_) return;
+  started_ = false;
+  nl_->WRCL().SetDefaultHandler(std::function<void(nl_msg*)>());
+
+  shm_xchg_->join();
+  shm_xchg_.reset();
+}
+
+void PacketSwitch::ProcessPacket(nl_msg* m, bool is_incoming) {
+  auto header = nlmsg_hdr(m);
+  auto genhdr = reinterpret_cast<genlmsghdr*>(nlmsg_data(header));
+
+  if (genhdr->cmd == WIFIROUTER_CMD_NOTIFY) {
+    // This attribute is mandatory: it contains MAC80211_HWSIM frame.
+    auto packet =
+        nlmsg_find_attr(header, sizeof(*genhdr), WIFIROUTER_ATTR_PACKET);
+    if (!packet) return;
+
+    // If origin is not local (= not set from local WIFI), then forward it to
+    // local WIFI.
+    if (is_incoming) {
+      // Need to update MAC80211_HWSIM WIFI family before injecting packet.
+      // Different kernels may have different family numbers allocated.
+      auto frame = reinterpret_cast<nlmsghdr*>(nla_data(packet));
+      frame->nlmsg_type = nl_->FamilyMAC80211();
+      frame->nlmsg_pid = 0;
+      frame->nlmsg_seq = 0;
+      frame->nlmsg_flags = NLM_F_REQUEST;
+      Cmd local(frame);
+
+      nl_->GeNL().Send(&local);
+
+      for (auto* r : local.Responses()) {
+        auto hdr = nlmsg_hdr(r);
+        if (hdr->nlmsg_type == NLMSG_ERROR) {
+          nlmsgerr* err = static_cast<nlmsgerr*>(nlmsg_data(hdr));
+          if (err->error < 0) {
+            LOG(ERROR) << "Could not send WIFI message: "
+                       << strerror(-err->error);
+          }
+        }
+      }
+    } else {
+      shm_wifi_.Send(header, header->nlmsg_len);
+    }
+  }
+}
+
+}  // namespace cvd
diff --git a/common/libs/wifi/packet_switch.h b/common/libs/wifi/packet_switch.h
new file mode 100644
index 0000000..0fbe213
--- /dev/null
+++ b/common/libs/wifi/packet_switch.h
@@ -0,0 +1,54 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <memory>
+#include <set>
+
+#include "common/libs/wifi/netlink.h"
+#include "common/libs/wifi/virtual_wifi.h"
+#include "common/libs/wifi/wr_client.h"
+#include "common/vsoc/lib/wifi_exchange_view.h"
+
+namespace cvd {
+
+class PacketSwitch {
+ public:
+  PacketSwitch(Netlink* nl) : nl_(nl) {}
+  ~PacketSwitch();
+
+  bool Init();
+  void Start();
+  void Stop();
+
+ private:
+  void ProcessPacket(nl_msg* m, bool is_incoming);
+
+  Netlink* nl_;
+
+  std::mutex op_mutex_;
+  // Started is referenced by all threads created by PacketSwitch to determine
+  // whether to carry on working, or terminate.
+  bool started_ = false;
+
+  std::unique_ptr<std::thread> shm_xchg_;
+  vsoc::wifi::WifiExchangeView shm_wifi_;
+
+  PacketSwitch(const PacketSwitch&) = delete;
+  PacketSwitch& operator=(const PacketSwitch&) = delete;
+};
+
+}  // namespace cvd
diff --git a/common/libs/wifi/router.h b/common/libs/wifi/router.h
index 9a5ab93..364de94 100644
--- a/common/libs/wifi/router.h
+++ b/common/libs/wifi/router.h
@@ -38,7 +38,7 @@
   WIFIROUTER_ATTR_UNSPEC,
 
   // MAC address representing interface from which the packet originated.
-  WIFIROUTER_ATTR_MAC,
+  WIFIROUTER_ATTR_HWSIM_ID,
 
   // MAC80211_HWSIM packet content.
   WIFIROUTER_ATTR_PACKET,
@@ -48,4 +48,3 @@
 };
 
 }  // namespace cvd
-
diff --git a/common/libs/wifi/virtual_wifi.cc b/common/libs/wifi/virtual_wifi.cc
index 365209e..140ff77 100644
--- a/common/libs/wifi/virtual_wifi.cc
+++ b/common/libs/wifi/virtual_wifi.cc
@@ -189,12 +189,12 @@
   return -1;
 }
 
-bool RegisterForRouterNotifications(Netlink* nl, uint8_t* mac_addr) {
+bool RegisterForRouterNotifications(Netlink* nl, int hwsim_id) {
   Cmd msg;
 
   if (!genlmsg_put(msg.Msg(), NL_AUTO_PID, NL_AUTO_SEQ, 0, 0,
                    NLM_F_REQUEST, WIFIROUTER_CMD_REGISTER, 0) ||
-      nla_put(msg.Msg(), WIFIROUTER_ATTR_MAC, MAX_ADDR_LEN, mac_addr)) {
+      nla_put_u32(msg.Msg(), WIFIROUTER_ATTR_HWSIM_ID, hwsim_id)) {
     LOG(ERROR) << "Could not create wifirouter register message.";
     return false;
   }
@@ -284,7 +284,7 @@
 
   // 5. Register with wifi router.
   LOG(INFO) << "Registering for notifications for: " << addr_;
-  if (!RegisterForRouterNotifications(nl_, mac_addr_)) {
+  if (!RegisterForRouterNotifications(nl_, hwsim_number_)) {
     LOG(ERROR) << "Could not register with wifi router.";
     return false;
   }
diff --git a/common/vsoc/lib/wifi_exchange_view.cpp b/common/vsoc/lib/wifi_exchange_view.cpp
index 23e0bc6..301f35f 100644
--- a/common/vsoc/lib/wifi_exchange_view.cpp
+++ b/common/vsoc/lib/wifi_exchange_view.cpp
@@ -20,17 +20,17 @@
 namespace vsoc {
 namespace wifi {
 
-bool WifiExchangeView::Send(const void* buffer, size_t length) {
+intptr_t WifiExchangeView::Send(const void* buffer, intptr_t length) {
 #ifdef CUTTLEFISH_HOST
   return data()->guest_ingress.Write(this, static_cast<const char*>(buffer),
-                                     length) == length;
+                                     length);
 #else
   return data()->guest_egress.Write(this, static_cast<const char*>(buffer),
-                                    length) == length;
+                                    length);
 #endif
 }
 
-intptr_t WifiExchangeView::Recv(void* buffer, size_t max_length) {
+intptr_t WifiExchangeView::Recv(void* buffer, intptr_t max_length) {
 #ifdef CUTTLEFISH_HOST
   return data()->guest_egress.Read(this, static_cast<char*>(buffer),
                                    max_length);
diff --git a/common/vsoc/lib/wifi_exchange_view.h b/common/vsoc/lib/wifi_exchange_view.h
index 83d7e6a..6bdab03 100644
--- a/common/vsoc/lib/wifi_exchange_view.h
+++ b/common/vsoc/lib/wifi_exchange_view.h
@@ -27,11 +27,11 @@
  public:
   // Send netlink packet to peer.
   // returns true, if operation was successful.
-  bool Send(const void* buffer, size_t length);
+  intptr_t Send(const void* buffer, intptr_t length);
 
   // Receive netlink packet from peer.
   // Returns number of bytes read, or negative value, if failed.
-  intptr_t Recv(void* buffer, size_t max_length);
+  intptr_t Recv(void* buffer, intptr_t max_length);
 
   // Set guest MAC address.
   void SetGuestMACAddress(const uint8_t* mac_address);
diff --git a/guest/commands/wifirouter/router.cc b/guest/commands/wifirouter/router.cc
index d01c3eb..a143ffb 100644
--- a/guest/commands/wifirouter/router.cc
+++ b/guest/commands/wifirouter/router.cc
@@ -37,10 +37,6 @@
 
 namespace cvd {
 namespace {
-using MacHash = uint64_t;
-using MacToClientsTable = std::multimap<MacHash, int>;
-using ClientsTable = std::set<int>;
-
 // Copied out of mac80211_hwsim.h header.
 constexpr int HWSIM_CMD_REGISTER = 1;
 constexpr int HWSIM_ATTR_ADDR_TRANSMITTER = 2;
@@ -50,13 +46,48 @@
 constexpr char kWifiSimFamilyName[] = "MAC80211_HWSIM";
 const int kMaxSupportedPacketSize = getpagesize();
 
-// Get hash for mac address serialized to 6 bytes of data starting at specified
-// location.
-// We don't care about byte ordering as much as we do about having all bytes
-// there. Byte order does not matter, we want to use it as a key in our map.
-uint64_t GetMacHash(const void* macaddr) {
-  auto typed = reinterpret_cast<const uint16_t*>(macaddr);
-  return (1ull * typed[0] << 32) | (typed[1] << 16) | typed[2];
+class WifiRouter {
+ public:
+  using MacHash = uint16_t;
+  using MacToClientsTable = std::multimap<MacHash, int>;
+  using ClientsTable = std::set<int>;
+
+  WifiRouter() : sock_(nullptr, nl_socket_free) {}
+  ~WifiRouter() = default;
+
+  void Init();
+  void ServerLoop();
+
+ private:
+  MacHash GetMacHash(const void* macaddr);
+  void CreateWifiRouterServerSocket();
+
+  void RegisterForHWSimNotifications();
+  void RouteWIFIPacket();
+
+  void AcceptNewClient();
+  bool HandleClientMessage(int client);
+  void RemoveClient(int client);
+
+  std::unique_ptr<nl_sock, void (*)(nl_sock*)> sock_;
+  int server_fd_ = 0;
+  int mac80211_family_ = 0;
+  ClientsTable registered_clients_;
+  MacToClientsTable registered_addresses_;
+};
+
+WifiRouter::MacHash WifiRouter::GetMacHash(const void* macaddr) {
+  const uint8_t* t = reinterpret_cast<const uint8_t*>(macaddr);
+
+  // This is guaranteed to be unique. Address here is assigned at creation time
+  // and is (well) non-mutable. This is a unique ID of the MAC80211 HWSIM
+  // interface.
+  return t[3] << 8 | t[4];
+}
+
+void WifiRouter::Init() {
+  CreateWifiRouterServerSocket();
+  RegisterForHWSimNotifications();
 }
 
 // Enable asynchronous notifications from MAC80211_HWSIM.
@@ -64,13 +95,29 @@
 // - `family` is MAC80211_HWSIM genl family number.
 //
 // Upon failure, this function will terminate execution of the program.
-void RegisterForHWSimNotifications(nl_sock* sock, int family) {
+void WifiRouter::RegisterForHWSimNotifications() {
+  sock_.reset(nl_socket_alloc());
+
+  auto res = nl_connect(sock_.get(), NETLINK_GENERIC);
+  if (res < 0) {
+    LOG(ERROR) << "Could not connect to netlink generic: " << nl_geterror(res);
+    exit(1);
+  }
+
+  mac80211_family_ = genl_ctrl_resolve(sock_.get(), kWifiSimFamilyName);
+  if (mac80211_family_ <= 0) {
+    LOG(ERROR) << "Could not find MAC80211 HWSIM. Please make sure module "
+               << "'mac80211_hwsim' is loaded on your system.";
+    exit(1);
+  }
+
   std::unique_ptr<nl_msg, void (*)(nl_msg*)> msg(
       nlmsg_alloc(), [](nl_msg* m) { nlmsg_free(m); });
-  genlmsg_put(msg.get(), NL_AUTO_PID, NL_AUTO_SEQ, family, 0, NLM_F_REQUEST,
-              HWSIM_CMD_REGISTER, 0);
-  nl_send_auto(sock, msg.get());
-  auto res = nl_wait_for_ack(sock);
+  genlmsg_put(msg.get(), NL_AUTO_PID, NL_AUTO_SEQ, mac80211_family_, 0,
+              NLM_F_REQUEST, HWSIM_CMD_REGISTER, 0);
+  nl_send_auto(sock_.get(), msg.get());
+
+  res = nl_wait_for_ack(sock_.get());
   if (res < 0) {
     LOG(ERROR) << "Could not register for notifications: " << nl_geterror(res);
     exit(1);
@@ -80,10 +127,10 @@
 // Create and configure WIFI Router server socket.
 // This function is guaranteed to success. If at any point an error is detected,
 // the function will terminate execution of the program.
-int CreateWifiRouterServerSocket() {
-  auto fd = socket(AF_UNIX, SOCK_SEQPACKET, 0);
-  if (fd <= 0) {
-    LOG(ERROR) << "Could not create unix socket: " << strerror(-fd);
+void WifiRouter::CreateWifiRouterServerSocket() {
+  server_fd_ = socket(AF_UNIX, SOCK_SEQPACKET, 0);
+  if (server_fd_ < 0) {
+    LOG(ERROR) << "Could not create unix socket: " << strerror(-errno);
     exit(1);
   }
 
@@ -92,39 +139,39 @@
   auto len = std::min(sizeof(addr.sun_path) - 2, FLAGS_socket_name.size());
   strncpy(&addr.sun_path[1], FLAGS_socket_name.c_str(), len);
   len += offsetof(sockaddr_un, sun_path) + 1;  // include heading \0 byte.
-  auto res = bind(fd, reinterpret_cast<sockaddr*>(&addr), len);
+  auto res = bind(server_fd_, reinterpret_cast<sockaddr*>(&addr), len);
 
   if (res < 0) {
-    LOG(ERROR) << "Could not bind unix socket: " << strerror(-res);
+    LOG(ERROR) << "Could not bind unix socket: " << strerror(-errno);
     exit(1);
   }
 
-  listen(fd, 4);
-  return fd;
+  listen(server_fd_, 4);
 }
 
 // Accept new WIFI Router client. When successful, client will be placed in
 // clients table.
-void AcceptNewClient(int server_fd, ClientsTable* clients) {
-  auto client = accept(server_fd, nullptr, nullptr);
+void WifiRouter::AcceptNewClient() {
+  auto client = accept(server_fd_, nullptr, nullptr);
   if (client < 0) {
-    LOG(ERROR) << "Could not accept client: " << strerror(errno);
+    LOG(ERROR) << "Could not accept client: " << strerror(-errno);
     return;
   }
 
-  clients->insert(client);
+  registered_clients_.insert(client);
   LOG(INFO) << "Client " << client << " added.";
 }
 
 // Disconnect and remove client from list of registered clients and recipients
 // of WLAN traffic.
-void RemoveClient(int client, ClientsTable* clients,
-                  MacToClientsTable* targets) {
+void WifiRouter::RemoveClient(int client) {
   close(client);
-  clients->erase(client);
-  for (auto iter = targets->begin(); iter != targets->end();) {
+  registered_clients_.erase(client);
+
+  for (auto iter = registered_addresses_.begin();
+       iter != registered_addresses_.end();) {
     if (iter->second == client) {
-      iter = targets->erase(iter);
+      iter = registered_addresses_.erase(iter);
     } else {
       ++iter;
     }
@@ -134,12 +181,11 @@
 
 // Read MAC80211HWSIM packet, find the originating MAC address and redirect it
 // to proper sink.
-void RouteWIFIPacket(nl_sock* nl, int simfamily, ClientsTable* clients,
-                     MacToClientsTable* targets) {
+void WifiRouter::RouteWIFIPacket() {
   sockaddr_nl tmp;
   uint8_t* buf;
 
-  const auto len = nl_recv(nl, &tmp, &buf, nullptr);
+  const auto len = nl_recv(sock_.get(), &tmp, &buf, nullptr);
   if (len < 0) {
     LOG(ERROR) << "Could not read from netlink: " << nl_geterror(len);
     return;
@@ -149,7 +195,7 @@
       reinterpret_cast<nlmsghdr*>(buf), [](nlmsghdr* m) { free(m); });
 
   // Discard messages that originate from anything else than MAC80211_HWSIM.
-  if (msg->nlmsg_type != simfamily) return;
+  if (msg->nlmsg_type != mac80211_family_) return;
 
   std::unique_ptr<nl_msg, void (*)(nl_msg*)> rep(
       nlmsg_alloc(), [](nl_msg* m) { nlmsg_free(m); });
@@ -163,28 +209,25 @@
   std::set<int> pending_removals;
   auto addr = attrs[HWSIM_ATTR_ADDR_TRANSMITTER];
   if (addr != nullptr) {
-    nla_put(rep.get(), WIFIROUTER_ATTR_MAC, nla_len(addr), nla_data(addr));
+    nla_put_u32(rep.get(), WIFIROUTER_ATTR_HWSIM_ID, GetMacHash(nla_data(addr)));
     nla_put(rep.get(), WIFIROUTER_ATTR_PACKET, len, buf);
     auto hdr = nlmsg_hdr(rep.get());
 
     auto key = GetMacHash(nla_data(attrs[HWSIM_ATTR_ADDR_TRANSMITTER]));
     LOG(INFO) << "Received netlink packet from " << std::hex << key;
-    for (auto it = targets->find(key); it != targets->end() && it->first == key;
-         ++it) {
-      auto num_written =
-          send(it->second, hdr, hdr->nlmsg_len, MSG_NOSIGNAL);
+    for (auto it = registered_addresses_.find(key);
+         it != registered_addresses_.end() && it->first == key; ++it) {
+      auto num_written = send(it->second, hdr, hdr->nlmsg_len, MSG_NOSIGNAL);
       if (num_written != static_cast<int64_t>(hdr->nlmsg_len)) {
         pending_removals.insert(it->second);
       }
     }
 
-    for (auto client : pending_removals) {
-      RemoveClient(client, clients, targets);
-    }
+    for (auto client : pending_removals) RemoveClient(client);
   }
 }
 
-bool HandleClientMessage(int client, MacToClientsTable* targets) {
+bool WifiRouter::HandleClientMessage(int client) {
   std::unique_ptr<nlmsghdr, void (*)(nlmsghdr*)> msg(
       reinterpret_cast<nlmsghdr*>(malloc(kMaxSupportedPacketSize)),
       [](nlmsghdr* h) { free(h); });
@@ -204,9 +247,15 @@
       nlattr* attrs[WIFIROUTER_ATTR_MAX];
       if (!nlmsg_parse(msg.get(), sizeof(genlmsghdr), attrs,
                        WIFIROUTER_ATTR_MAX - 1, nullptr)) {
-        if (attrs[WIFIROUTER_ATTR_MAC] != nullptr) {
-          targets->emplace(GetMacHash(nla_data(attrs[WIFIROUTER_ATTR_MAC])),
-                           client);
+        if (attrs[WIFIROUTER_ATTR_HWSIM_ID] != nullptr) {
+          LOG(INFO) << "Registering new client to receive data for "
+                    << nla_get_u32(attrs[WIFIROUTER_ATTR_HWSIM_ID]);
+          registered_addresses_.emplace(
+              nla_get_u32(attrs[WIFIROUTER_ATTR_HWSIM_ID]), client);
+          // This is unfortunate, but it is a bug in mac80211_hwsim stack.
+          // Apparently, the imperfect medium will not receive notifications for
+          // newly created wifi interfaces. How about that...
+          RegisterForHWSimNotifications();
           result = 0;
         }
       }
@@ -229,13 +278,9 @@
 }
 
 // Process incoming requests from netlink, server or clients.
-void ServerLoop(int server_fd, nl_sock* netlink_sock, int family) {
-  ClientsTable clients;
-  MacToClientsTable targets;
-  int netlink_fd = nl_socket_get_fd(netlink_sock);
-
+void WifiRouter::ServerLoop() {
   while (true) {
-    auto max_fd = server_fd;
+    auto max_fd = 0;
     fd_set reads{};
 
     auto fdset = [&max_fd, &reads](int fd) {
@@ -243,26 +288,23 @@
       max_fd = std::max(max_fd, fd);
     };
 
-    fdset(server_fd);
-    fdset(netlink_fd);
-    for (int client : clients) fdset(client);
+    fdset(server_fd_);
+    fdset(nl_socket_get_fd(sock_.get()));
+    for (int client : registered_clients_) fdset(client);
 
     if (select(max_fd + 1, &reads, nullptr, nullptr, nullptr) <= 0) continue;
 
-    if (FD_ISSET(server_fd, &reads)) AcceptNewClient(server_fd, &clients);
-    if (FD_ISSET(netlink_fd, &reads))
-      RouteWIFIPacket(netlink_sock, family, &clients, &targets);
+    if (FD_ISSET(server_fd_, &reads)) AcceptNewClient();
+    if (FD_ISSET(nl_socket_get_fd(sock_.get()), &reads)) RouteWIFIPacket();
 
     // Process any client messages left. Drop any client that is no longer
     // talking with us.
-    for (auto client = clients.begin(); client != clients.end();) {
+    for (auto client = registered_clients_.begin();
+         client != registered_clients_.end();) {
       auto cfd = *client++;
       // Is our client sending us data?
       if (FD_ISSET(cfd, &reads)) {
-        if (!HandleClientMessage(cfd, &targets)) {
-          // Client should be disconnected.
-          RemoveClient(cfd, &clients, &targets);
-        }
+        if (!HandleClientMessage(cfd)) RemoveClient(cfd);
       }
     }
   }
@@ -280,23 +322,7 @@
   google::InstallFailureSignalHandler();
 #endif
 
-  std::unique_ptr<nl_sock, void (*)(nl_sock*)> sock(nl_socket_alloc(),
-                                                    nl_socket_free);
-
-  auto res = nl_connect(sock.get(), NETLINK_GENERIC);
-  if (res < 0) {
-    LOG(ERROR) << "Could not connect to netlink generic: " << nl_geterror(res);
-    exit(1);
-  }
-
-  auto mac80211_family = genl_ctrl_resolve(sock.get(), kWifiSimFamilyName);
-  if (mac80211_family <= 0) {
-    LOG(ERROR) << "Could not find MAC80211 HWSIM. Please make sure module "
-               << "'mac80211_hwsim' is loaded on your system.";
-    exit(1);
-  }
-
-  RegisterForHWSimNotifications(sock.get(), mac80211_family);
-  auto server_fd = CreateWifiRouterServerSocket();
-  ServerLoop(server_fd, sock.get(), mac80211_family);
+  WifiRouter r;
+  r.Init();
+  r.ServerLoop();
 }