From dd495a9da8cd7b4ffc7071bb43a67be948c2ca96 Mon Sep 17 00:00:00 2001
From: Natale Patriciello <natale.patriciello@gmail.com>
Date: Fri, 8 Mar 2019 15:41:59 +0100
Subject: [PATCH] lte: Make point-to-point-epc-helper inheritable

---
 src/lte/helper/point-to-point-epc-helper.cc | 47 ++++++++++++++++-----
 src/lte/helper/point-to-point-epc-helper.h  | 25 +++++++++++
 2 files changed, 62 insertions(+), 10 deletions(-)

diff --git a/src/lte/helper/point-to-point-epc-helper.cc b/src/lte/helper/point-to-point-epc-helper.cc
index fe6304e65..b60c87a10 100644
--- a/src/lte/helper/point-to-point-epc-helper.cc
+++ b/src/lte/helper/point-to-point-epc-helper.cc
@@ -477,20 +477,35 @@ PointToPointEpcHelper::AddX2Interface (Ptr<Node> enb1, Ptr<Node> enb2)
 
   // Add X2 interface to both eNBs' X2 entities
   Ptr<EpcX2> enb1X2 = enb1->GetObject<EpcX2> ();
-  Ptr<LteEnbNetDevice> enb1LteDev = enb1->GetDevice (0)->GetObject<LteEnbNetDevice> ();
-  uint16_t enb1CellId = enb1LteDev->GetCellId ();
-  NS_LOG_LOGIC ("LteEnbNetDevice #1 = " << enb1LteDev << " - CellId = " << enb1CellId);
-
   Ptr<EpcX2> enb2X2 = enb2->GetObject<EpcX2> ();
-  Ptr<LteEnbNetDevice> enb2LteDev = enb2->GetDevice (0)->GetObject<LteEnbNetDevice> ();
-  uint16_t enb2CellId = enb2LteDev->GetCellId ();
+
+  Ptr<NetDevice> enb1LteDev = enb1->GetDevice (0); // What if at position 0 we have another device than LTE?
+  Ptr<NetDevice> enb2LteDev = enb2->GetDevice (0); // What if at position 0 we have another device than LTE?
+
+  DoAddX2Interface (enb1X2, enb1LteDev, enb1X2Address, enb2X2, enb2LteDev, enb2X2Address);
+}
+
+void
+PointToPointEpcHelper::DoAddX2Interface (const Ptr<EpcX2> &enb1X2, const Ptr<NetDevice> &enb1LteDev,
+                                         const Ipv4Address &enb1X2Address,
+                                         const Ptr<EpcX2> &enb2X2, const Ptr<NetDevice> &enb2LteDev,
+                                         const Ipv4Address &enb2X2Address) const
+{
+  NS_LOG_FUNCTION (this);
+
+  Ptr<LteEnbNetDevice> enb1LteDevice = enb1LteDev->GetObject<LteEnbNetDevice> ();
+  Ptr<LteEnbNetDevice> enb2LteDevice = enb2LteDev->GetObject<LteEnbNetDevice> ();
+  uint16_t enb1CellId = enb1LteDevice->GetCellId ();
+  uint16_t enb2CellId = enb2LteDevice->GetCellId ();
+
+  NS_LOG_LOGIC ("LteEnbNetDevice #1 = " << enb1LteDev << " - CellId = " << enb1CellId);
   NS_LOG_LOGIC ("LteEnbNetDevice #2 = " << enb2LteDev << " - CellId = " << enb2CellId);
 
   enb1X2->AddX2Interface (enb1CellId, enb1X2Address, enb2CellId, enb2X2Address);
   enb2X2->AddX2Interface (enb2CellId, enb2X2Address, enb1CellId, enb1X2Address);
 
-  enb1LteDev->GetRrc ()->AddX2Neighbour (enb2LteDev->GetCellId ());
-  enb2LteDev->GetRrc ()->AddX2Neighbour (enb1LteDev->GetCellId ());
+  enb1LteDevice->GetRrc ()->AddX2Neighbour (enb2CellId);
+  enb2LteDevice->GetRrc ()->AddX2Neighbour (enb1CellId);
 }
 
 
@@ -504,7 +519,8 @@ PointToPointEpcHelper::AddUe (Ptr<NetDevice> ueDevice, uint64_t imsi)
 }
 
 uint8_t
-PointToPointEpcHelper::ActivateEpsBearer (Ptr<NetDevice> ueDevice, uint64_t imsi, Ptr<EpcTft> tft, EpsBearer bearer)
+PointToPointEpcHelper::ActivateEpsBearer (Ptr<NetDevice> ueDevice, uint64_t imsi,
+                                          Ptr<EpcTft> tft, EpsBearer bearer)
 {
   NS_LOG_FUNCTION (this << ueDevice << imsi);
 
@@ -537,12 +553,23 @@ PointToPointEpcHelper::ActivateEpsBearer (Ptr<NetDevice> ueDevice, uint64_t imsi
         }
     }
   uint8_t bearerId = m_mmeApp->AddBearer (imsi, tft, bearer);
+
+  DoActivateEpsBearerForUe (ueDevice, tft, bearer);
+
+  return bearerId;
+}
+
+void
+PointToPointEpcHelper::DoActivateEpsBearerForUe (const Ptr<NetDevice> &ueDevice,
+                                                 const Ptr<EpcTft> &tft,
+                                                 const EpsBearer &bearer) const
+{
+  NS_LOG_FUNCTION (this);
   Ptr<LteUeNetDevice> ueLteDevice = ueDevice->GetObject<LteUeNetDevice> ();
   if (ueLteDevice)
     {
       Simulator::ScheduleNow (&EpcUeNas::ActivateEpsBearer, ueLteDevice->GetNas (), bearer, tft);
     }
-  return bearerId;
 }
 
 Ptr<Node>
diff --git a/src/lte/helper/point-to-point-epc-helper.h b/src/lte/helper/point-to-point-epc-helper.h
index b2e252aa5..9d60c2513 100644
--- a/src/lte/helper/point-to-point-epc-helper.h
+++ b/src/lte/helper/point-to-point-epc-helper.h
@@ -83,6 +83,31 @@ public:
   virtual Ipv4Address GetUeDefaultGatewayAddress ();
   virtual Ipv6Address GetUeDefaultGatewayAddress6 ();
 
+protected:
+  /**
+   * \brief DoAddX2Interface: Call AddX2Interface on top of the Enb device pointers
+   *
+   * \param enb1X2 EPCX2 of ENB1
+   * \param enb1LteDev LTE device of ENB1
+   * \param enb1X2Address Address for ENB1
+   * \param enb2X2 EPCX2 of ENB2
+   * \param enb2LteDev LTE device of ENB2
+   * \param enb2X2Address Address for ENB2
+   */
+  virtual void DoAddX2Interface(const Ptr<EpcX2> &enb1X2, const Ptr<NetDevice> &enb1LteDev,
+                                const Ipv4Address &enb1X2Address,
+                                const Ptr<EpcX2> &enb2X2, const Ptr<NetDevice> &enb2LteDev,
+                                const Ipv4Address &enb2X2Address) const;
+
+  /**
+   * \brief DoActivateEpsBearerForUe: Schedule ActivateEpsBearer on the UE
+   * \param ueDevice LTE device for the UE
+   * \param tft TFT
+   * \param bearer Bearer
+   */
+  virtual void DoActivateEpsBearerForUe (const Ptr<NetDevice> &ueDevice,
+                                         const Ptr<EpcTft> &tft,
+                                         const EpsBearer &bearer) const;
 
 private: