From b921650795f616b7a63b138615c63b51f32d2968 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Deronne?= <sebastien.deronne@gmail.com>
Date: Tue, 2 Jan 2024 10:52:56 +0100
Subject: [PATCH] wifi: Cleanup duplicated code in WifiNetDevice

---
 src/wifi/model/wifi-net-device.cc | 45 ++++++++++++++++++++-----------
 src/wifi/model/wifi-net-device.h  | 15 +++++++++++
 2 files changed, 44 insertions(+), 16 deletions(-)

diff --git a/src/wifi/model/wifi-net-device.cc b/src/wifi/model/wifi-net-device.cc
index 08dbc9dc2..d59ed8e97 100644
--- a/src/wifi/model/wifi-net-device.cc
+++ b/src/wifi/model/wifi-net-device.cc
@@ -478,17 +478,7 @@ bool
 WifiNetDevice::Send(Ptr<Packet> packet, const Address& dest, uint16_t protocolNumber)
 {
     NS_LOG_FUNCTION(this << packet << dest << protocolNumber);
-    NS_ASSERT(Mac48Address::IsMatchingType(dest));
-
-    Mac48Address realTo = Mac48Address::ConvertFrom(dest);
-
-    LlcSnapHeader llc;
-    llc.SetType(protocolNumber);
-    packet->AddHeader(llc);
-
-    m_mac->NotifyTx(packet);
-    m_mac->Enqueue(packet, realTo);
-    return true;
+    return DoSend(packet, std::nullopt, dest, protocolNumber);
 }
 
 Ptr<Node>
@@ -579,18 +569,41 @@ WifiNetDevice::SendFrom(Ptr<Packet> packet,
                         uint16_t protocolNumber)
 {
     NS_LOG_FUNCTION(this << packet << source << dest << protocolNumber);
-    NS_ASSERT(Mac48Address::IsMatchingType(dest));
-    NS_ASSERT(Mac48Address::IsMatchingType(source));
+    return DoSend(packet, source, dest, protocolNumber);
+}
 
-    Mac48Address realTo = Mac48Address::ConvertFrom(dest);
-    Mac48Address realFrom = Mac48Address::ConvertFrom(source);
+bool
+WifiNetDevice::DoSend(Ptr<Packet> packet,
+                      std::optional<Address> source,
+                      const Address& dest,
+                      uint16_t protocolNumber)
+{
+    NS_LOG_FUNCTION(this << packet << dest << protocolNumber << source.value_or(Address()));
+
+    if (source)
+    {
+        NS_ASSERT_MSG(Mac48Address::IsMatchingType(*source),
+                      *source << " is not compatible with a Mac48Address");
+    }
+    NS_ASSERT_MSG(Mac48Address::IsMatchingType(dest),
+                  dest << " is not compatible with a Mac48Address");
+
+    auto realTo = Mac48Address::ConvertFrom(dest);
 
     LlcSnapHeader llc;
     llc.SetType(protocolNumber);
     packet->AddHeader(llc);
 
     m_mac->NotifyTx(packet);
-    m_mac->Enqueue(packet, realTo, realFrom);
+    if (source)
+    {
+        auto realFrom = Mac48Address::ConvertFrom(*source);
+        m_mac->Enqueue(packet, realTo, realFrom);
+    }
+    else
+    {
+        m_mac->Enqueue(packet, realTo);
+    }
 
     return true;
 }
diff --git a/src/wifi/model/wifi-net-device.h b/src/wifi/model/wifi-net-device.h
index 81901da09..c1ab4994a 100644
--- a/src/wifi/model/wifi-net-device.h
+++ b/src/wifi/model/wifi-net-device.h
@@ -25,6 +25,7 @@
 #include "ns3/net-device.h"
 #include "ns3/traced-callback.h"
 
+#include <optional>
 #include <vector>
 
 namespace ns3
@@ -239,6 +240,20 @@ class WifiNetDevice : public NetDevice
      */
     void CompleteConfig();
 
+    /**
+     * Send a packet
+     *
+     * \param packet packet sent from upper layers
+     * \param source source mac address (if provided)
+     * \param dest destination mac address
+     * \param protocolNumber type of payload contained in this packet
+     * \returns true if successful, false otherwise
+     */
+    bool DoSend(Ptr<Packet> packet,
+                std::optional<Address> source,
+                const Address& dest,
+                uint16_t protocolNumber);
+
     Ptr<Node> m_node;                                             //!< the node
     std::vector<Ptr<WifiPhy>> m_phys;                             //!< the phy objects
     Ptr<WifiMac> m_mac;                                           //!< the MAC