3 files changed, 250 insertions, 19 deletions

diff --git a/src/core/hle/service/nwm/nwm_uds.cpp b/src/core/hle/service/nwm/nwm_uds.cpp
index 8ef0cda09..0aa63cc1e 100644
--- a/src/core/hle/service/nwm/nwm_uds.cpp
+++ b/src/core/hle/service/nwm/nwm_uds.cpp
@@ -15,6 +15,7 @@
 #include "core/hle/ipc_helpers.h"
 #include "core/hle/kernel/event.h"
 #include "core/hle/kernel/shared_memory.h"
+#include "core/hle/lock.h"
 #include "core/hle/result.h"
 #include "core/hle/service/nwm/nwm_uds.h"
 #include "core/hle/service/nwm/uds_beacon.h"
@@ -100,6 +101,20 @@ void SendPacket(Network::WifiPacket& packet) {
     // TODO(Subv): Implement.
 }
 
+/*
+ * Returns an available index in the nodes array for the
+ * currently-hosted UDS network.
+ */
+static u16 GetNextAvailableNodeId() {
+    for (u16 index = 0; index < connection_status.max_nodes; ++index) {
+        if ((connection_status.node_bitmask & (1 << index)) == 0)
+            return index;
+    }
+
+    // Any connection attempts to an already full network should have been refused.
+    ASSERT_MSG(false, "No available connection slots in the network");
+}
+
 // Inserts the received beacon frame in the beacon queue and removes any older beacons if the size
 // limit is exceeded.
 void HandleBeaconFrame(const Network::WifiPacket& packet) {
@@ -143,18 +158,88 @@ void HandleAssociationResponseFrame(const Network::WifiPacket& packet) {
     SendPacket(eapol_start);
 }
 
-/*
- * Returns an available index in the nodes array for the
- * currently-hosted UDS network.
- */
-static u16 GetNextAvailableNodeId() {
-    for (u16 index = 0; index < connection_status.max_nodes; ++index) {
-        if ((connection_status.node_bitmask & (1 << index)) == 0)
-            return index;
-    }
+static void HandleEAPoLPacket(const Network::WifiPacket& packet) {
+    std::lock_guard<std::mutex> lock(connection_status_mutex);
 
-    // Any connection attempts to an already full network should have been refused.
-    ASSERT_MSG(false, "No available connection slots in the network");
+    if (GetEAPoLFrameType(packet.data) == EAPoLStartMagic) {
+        if (connection_status.status != static_cast<u32>(NetworkStatus::ConnectedAsHost)) {
+            LOG_DEBUG(Service_NWM, "Connection sequence aborted, because connection status is %u",
+                      connection_status.status);
+            return;
+        }
+
+        auto node = DeserializeNodeInfoFromFrame(packet.data);
+
+        if (connection_status.max_nodes == connection_status.total_nodes) {
+            // Reject connection attempt
+            LOG_ERROR(Service_NWM, "Reached maximum nodes, but reject packet wasn't sent.");
+            // TODO(B3N30): Figure out what packet is sent here
+            return;
+        }
+
+        // Get an unused network node id
+        u16 node_id = GetNextAvailableNodeId();
+        node.network_node_id = node_id + 1;
+
+        connection_status.node_bitmask |= 1 << node_id;
+        connection_status.changed_nodes |= 1 << node_id;
+        connection_status.nodes[node_id] = node.network_node_id;
+        connection_status.total_nodes++;
+
+        u8 current_nodes = network_info.total_nodes;
+        node_info[current_nodes] = node;
+
+        network_info.total_nodes++;
+
+        // Send the EAPoL-Logoff packet.
+        using Network::WifiPacket;
+        WifiPacket eapol_logoff;
+        eapol_logoff.channel = network_channel;
+        eapol_logoff.data =
+            GenerateEAPoLLogoffFrame(packet.transmitter_address, node.network_node_id, node_info,
+                                     network_info.max_nodes, network_info.total_nodes);
+        // TODO(Subv): Encrypt the packet.
+        eapol_logoff.destination_address = packet.transmitter_address;
+        eapol_logoff.type = WifiPacket::PacketType::Data;
+
+        SendPacket(eapol_logoff);
+        // TODO(B3N30): Broadcast updated node list
+        // The 3ds does this presumably to support spectators.
+        std::lock_guard<std::recursive_mutex> lock(HLE::g_hle_lock);
+        connection_status_event->Signal();
+    } else {
+        if (connection_status.status != static_cast<u32>(NetworkStatus::NotConnected)) {
+            LOG_DEBUG(Service_NWM, "Connection sequence aborted, because connection status is %u",
+                      connection_status.status);
+            return;
+        }
+        auto logoff = ParseEAPoLLogoffFrame(packet.data);
+
+        network_info.total_nodes = logoff.connected_nodes;
+        network_info.max_nodes = logoff.max_nodes;
+
+        connection_status.network_node_id = logoff.assigned_node_id;
+        connection_status.total_nodes = logoff.connected_nodes;
+        connection_status.max_nodes = logoff.max_nodes;
+
+        node_info.clear();
+        node_info.reserve(network_info.max_nodes);
+        for (size_t index = 0; index < logoff.connected_nodes; ++index) {
+            connection_status.node_bitmask |= 1 << index;
+            connection_status.changed_nodes |= 1 << index;
+            connection_status.nodes[index] = logoff.nodes[index].network_node_id;
+
+            node_info.emplace_back(DeserializeNodeInfo(logoff.nodes[index]));
+        }
+
+        // We're now connected, signal the application
+        connection_status.status = static_cast<u32>(NetworkStatus::ConnectedAsClient);
+        // Some games require ConnectToNetwork to block, for now it doesn't
+        // If blocking is implemented this lock needs to be changed,
+        // otherwise it might cause deadlocks
+        std::lock_guard<std::recursive_mutex> lock(HLE::g_hle_lock);
+        connection_status_event->Signal();
+    }
 }
 
 /*
@@ -238,6 +323,17 @@ void HandleAuthenticationFrame(const Network::WifiPacket& packet) {
     }
 }
 
+static void HandleDataFrame(const Network::WifiPacket& packet) {
+    switch (GetFrameEtherType(packet.data)) {
+    case EtherType::EAPoL:
+        HandleEAPoLPacket(packet);
+        break;
+    case EtherType::SecureData:
+        // TODO(B3N30): Handle SecureData packets
+        break;
+    }
+}
+
 /// Callback to parse and handle a received wifi packet.
 void OnWifiPacketReceived(const Network::WifiPacket& packet) {
     switch (packet.type) {
@@ -250,6 +346,9 @@ void OnWifiPacketReceived(const Network::WifiPacket& packet) {
     case Network::WifiPacket::PacketType::AssociationResponse:
         HandleAssociationResponseFrame(packet);
         break;
+    case Network::WifiPacket::PacketType::Data:
+        HandleDataFrame(packet);
+        break;
     }
 }
 
diff --git a/src/core/hle/service/nwm/uds_data.cpp b/src/core/hle/service/nwm/uds_data.cpp
index 3ef2a84b6..4b389710f 100644
--- a/src/core/hle/service/nwm/uds_data.cpp
+++ b/src/core/hle/service/nwm/uds_data.cpp
@@ -2,6 +2,7 @@
 // Licensed under GPLv2 or any later version
 // Refer to the license.txt file included.
 
+#include <algorithm>
 #include <cstring>
 #include <cryptopp/aes.h>
 #include <cryptopp/ccm.h>
@@ -277,10 +278,10 @@ std::vector<u8> GenerateDataPayload(const std::vector<u8>& data, u8 channel, u16
 std::vector<u8> GenerateEAPoLStartFrame(u16 association_id, const NodeInfo& node_info) {
     EAPoLStartPacket eapol_start{};
     eapol_start.association_id = association_id;
-    eapol_start.friend_code_seed = node_info.friend_code_seed;
+    eapol_start.node.friend_code_seed = node_info.friend_code_seed;
 
-    for (int i = 0; i < node_info.username.size(); ++i)
-        eapol_start.username[i] = node_info.username[i];
+    std::copy(node_info.username.begin(), node_info.username.end(),
+              eapol_start.node.username.begin());
 
     // Note: The network_node_id and unknown bytes seem to be uninitialized in the NWM module.
     // TODO(B3N30): The last 8 bytes seem to have a fixed value of 07 88 15 00 04 e9 13 00 in
@@ -295,5 +296,78 @@ std::vector<u8> GenerateEAPoLStartFrame(u16 association_id, const NodeInfo& node
     return buffer;
 }
 
+EtherType GetFrameEtherType(const std::vector<u8>& frame) {
+    LLCHeader header;
+    std::memcpy(&header, frame.data(), sizeof(header));
+
+    u16 ethertype = header.protocol;
+    return static_cast<EtherType>(ethertype);
+}
+
+u16 GetEAPoLFrameType(const std::vector<u8>& frame) {
+    // Ignore the LLC header
+    u16_be eapol_type;
+    std::memcpy(&eapol_type, frame.data() + sizeof(LLCHeader), sizeof(eapol_type));
+    return eapol_type;
+}
+
+NodeInfo DeserializeNodeInfoFromFrame(const std::vector<u8>& frame) {
+    EAPoLStartPacket eapol_start;
+
+    // Skip the LLC header
+    std::memcpy(&eapol_start, frame.data() + sizeof(LLCHeader), sizeof(eapol_start));
+
+    NodeInfo node{};
+    node.friend_code_seed = eapol_start.node.friend_code_seed;
+
+    std::copy(eapol_start.node.username.begin(), eapol_start.node.username.end(),
+              node.username.begin());
+
+    return node;
+}
+
+NodeInfo DeserializeNodeInfo(const EAPoLNodeInfo& node) {
+    NodeInfo node_info{};
+    node_info.friend_code_seed = node.friend_code_seed;
+    node_info.network_node_id = node.network_node_id;
+
+    std::copy(node.username.begin(), node.username.end(), node_info.username.begin());
+
+    return node_info;
+}
+
+std::vector<u8> GenerateEAPoLLogoffFrame(const MacAddress& mac_address, u16 network_node_id,
+                                         const NodeList& nodes, u8 max_nodes, u8 total_nodes) {
+    EAPoLLogoffPacket eapol_logoff{};
+    eapol_logoff.assigned_node_id = network_node_id;
+    eapol_logoff.connected_nodes = total_nodes;
+    eapol_logoff.max_nodes = max_nodes;
+
+    for (size_t index = 0; index < total_nodes; ++index) {
+        const auto& node_info = nodes[index];
+        auto& node = eapol_logoff.nodes[index];
+
+        node.friend_code_seed = node_info.friend_code_seed;
+        node.network_node_id = node_info.network_node_id;
+
+        std::copy(node_info.username.begin(), node_info.username.end(), node.username.begin());
+    }
+
+    std::vector<u8> eapol_buffer(sizeof(EAPoLLogoffPacket));
+    std::memcpy(eapol_buffer.data(), &eapol_logoff, sizeof(eapol_logoff));
+
+    std::vector<u8> buffer = GenerateLLCHeader(EtherType::EAPoL);
+    buffer.insert(buffer.end(), eapol_buffer.begin(), eapol_buffer.end());
+    return buffer;
+}
+
+EAPoLLogoffPacket ParseEAPoLLogoffFrame(const std::vector<u8>& frame) {
+    EAPoLLogoffPacket eapol_logoff;
+
+    // Skip the LLC header
+    std::memcpy(&eapol_logoff, frame.data() + sizeof(LLCHeader), sizeof(eapol_logoff));
+    return eapol_logoff;
+}
+
 } // namespace NWM
 } // namespace Service
diff --git a/src/core/hle/service/nwm/uds_data.h b/src/core/hle/service/nwm/uds_data.h
index 76e8f546b..76bccb1bf 100644
--- a/src/core/hle/service/nwm/uds_data.h
+++ b/src/core/hle/service/nwm/uds_data.h
@@ -8,6 +8,7 @@
 #include <vector>
 #include "common/common_types.h"
 #include "common/swap.h"
+#include "core/hle/service/nwm/uds_beacon.h"
 #include "core/hle/service/service.h"
 
 namespace Service {
@@ -67,6 +68,16 @@ struct DataFrameCryptoCTR {
 
 static_assert(sizeof(DataFrameCryptoCTR) == 16, "DataFrameCryptoCTR has the wrong size");
 
+struct EAPoLNodeInfo {
+    u64_be friend_code_seed;
+    std::array<u16_be, 10> username;
+    INSERT_PADDING_BYTES(4);
+    u16_be network_node_id;
+    INSERT_PADDING_BYTES(6);
+};
+
+static_assert(sizeof(EAPoLNodeInfo) == 0x28, "EAPoLNodeInfo has the wrong size");
+
 constexpr u16 EAPoLStartMagic = 0x201;
 
 /*
@@ -78,15 +89,27 @@ struct EAPoLStartPacket {
     // This value is hardcoded to 1 in the NWM module.
     u16_be unknown = 1;
     INSERT_PADDING_BYTES(2);
+    EAPoLNodeInfo node;
+};
 
-    u64_be friend_code_seed;
-    std::array<u16_be, 10> username;
-    INSERT_PADDING_BYTES(4);
-    u16_be network_node_id;
+static_assert(sizeof(EAPoLStartPacket) == 0x30, "EAPoLStartPacket has the wrong size");
+
+constexpr u16 EAPoLLogoffMagic = 0x202;
+
+struct EAPoLLogoffPacket {
+    u16_be magic = EAPoLLogoffMagic;
+    INSERT_PADDING_BYTES(2);
+    u16_be assigned_node_id;
+    MacAddress client_mac_address;
     INSERT_PADDING_BYTES(6);
+    u8 connected_nodes;
+    u8 max_nodes;
+    INSERT_PADDING_BYTES(4);
+
+    std::array<EAPoLNodeInfo, UDSMaxNodes> nodes;
 };
 
-static_assert(sizeof(EAPoLStartPacket) == 0x30, "EAPoLStartPacket has the wrong size");
+static_assert(sizeof(EAPoLLogoffPacket) == 0x298, "EAPoLLogoffPacket has the wrong size");
 
 /**
  * Generates an unencrypted 802.11 data payload.
@@ -102,5 +125,40 @@ std::vector<u8> GenerateDataPayload(const std::vector<u8>& data, u8 channel, u16
  */
 std::vector<u8> GenerateEAPoLStartFrame(u16 association_id, const NodeInfo& node_info);
 
+/*
+ * Returns the EtherType of the specified 802.11 frame.
+ */
+EtherType GetFrameEtherType(const std::vector<u8>& frame);
+
+/*
+ * Returns the EAPoL type (Start / Logoff) of the specified 802.11 frame.
+ * Note: The frame *must* be an EAPoL frame.
+ */
+u16 GetEAPoLFrameType(const std::vector<u8>& frame);
+
+/*
+ * Returns a deserialized NodeInfo structure from the information inside an EAPoL-Start packet
+ * encapsulated in an 802.11 data frame.
+ */
+NodeInfo DeserializeNodeInfoFromFrame(const std::vector<u8>& frame);
+
+/*
+ * Returns a NodeInfo constructed from the data in the specified EAPoLNodeInfo.
+ */
+NodeInfo DeserializeNodeInfo(const EAPoLNodeInfo& node);
+
+/*
+ * Generates an unencrypted 802.11 data frame body with the EAPoL-Logoff format for UDS
+ * communication.
+ * @returns The generated frame body.
+ */
+std::vector<u8> GenerateEAPoLLogoffFrame(const MacAddress& mac_address, u16 network_node_id,
+                                         const NodeList& nodes, u8 max_nodes, u8 total_nodes);
+
+/*
+ * Returns a EAPoLLogoffPacket representing the specified 802.11-encapsulated data frame.
+ */
+EAPoLLogoffPacket ParseEAPoLLogoffFrame(const std::vector<u8>& frame);
+
 } // namespace NWM
 } // namespace Service