From ed946b1172ed2eaf73b68d8d08918ccde0ea5d2a Mon Sep 17 00:00:00 2001
From: aligungr <aligng1620@gmail.com>
Date: Thu, 18 Feb 2021 00:53:22 +0300
Subject: [PATCH] NNSF refactor

---
 src/gnb/ngap/management.cpp | 19 ++++++-------------
 src/gnb/ngap/nas.cpp        |  2 +-
 src/gnb/ngap/nnsf.cpp       | 28 ++++++++++++++++++++++++++++
 src/gnb/ngap/task.hpp       |  5 ++++-
 4 files changed, 39 insertions(+), 15 deletions(-)
 create mode 100644 src/gnb/ngap/nnsf.cpp

diff --git a/src/gnb/ngap/management.cpp b/src/gnb/ngap/management.cpp
index 93c7e08f0..edf508de7 100644
--- a/src/gnb/ngap/management.cpp
+++ b/src/gnb/ngap/management.cpp
@@ -40,13 +40,12 @@ void NgapTask::createUeContext(int ueId)
 
     m_ueCtx[ctx->ctxId] = ctx;
 
-    // Select an AMF for the UE (if any)
-    for (auto &amf : m_amfCtx)
-    {
-        // todo: an arbitrary AMF is selected for now
-        ctx->associatedAmfId = amf.second->ctxId;
-        break;
-    }
+    // Perform AMF selection
+    auto *amf = selectAmf(ueId);
+    if (amf == nullptr)
+        m_logger->err("AMF selection for UE[%d] failed. Could not find a suitable AMF.", ueId);
+    else
+        ctx->associatedAmfId = amf->ctxId;
 }
 
 NgapUeContext *NgapTask::findUeContext(int ctxId)
@@ -134,12 +133,6 @@ NgapUeContext *NgapTask::findUeByNgapIdPair(int amfCtxId, const NgapIdPair &idPa
     return ue;
 }
 
-NgapAmfContext *NgapTask::selectNewAmfForReAllocation(int initiatedAmfId, int amfSetId)
-{
-    // TODO an arbitrary AMF is selected for now
-    return findAmfContext(initiatedAmfId);
-}
-
 void NgapTask::deleteUeContext(int ueId)
 {
     auto *ue = m_ueCtx[ueId];
diff --git a/src/gnb/ngap/nas.cpp b/src/gnb/ngap/nas.cpp
index 23d348a41..5db9a9d16 100644
--- a/src/gnb/ngap/nas.cpp
+++ b/src/gnb/ngap/nas.cpp
@@ -174,7 +174,7 @@ void NgapTask::receiveRerouteNasRequest(int amfId, ASN_NGAP_RerouteNASRequest *m
         }
     }
 
-    auto *newAmf = selectNewAmfForReAllocation(amfId, asn::GetBitStringInt<10>(ieAmfSetId->AMFSetID));
+    auto *newAmf = selectNewAmfForReAllocation(ue->ctxId, amfId, asn::GetBitStringInt<10>(ieAmfSetId->AMFSetID));
     if (newAmf == nullptr)
     {
         m_logger->err("AMF selection for re-allocation failed. Could not find a suitable AMF.");
diff --git a/src/gnb/ngap/nnsf.cpp b/src/gnb/ngap/nnsf.cpp
new file mode 100644
index 000000000..4c22e91ac
--- /dev/null
+++ b/src/gnb/ngap/nnsf.cpp
@@ -0,0 +1,28 @@
+//
+// This file is a part of UERANSIM open source project.
+// Copyright (c) 2021 ALİ GÜNGÖR.
+//
+// The software and all associated files are licensed under GPL-3.0
+// and subject to the terms and conditions defined in LICENSE file.
+//
+
+#include "task.hpp"
+
+namespace nr::gnb
+{
+
+NgapAmfContext *NgapTask::selectAmf(int ueId)
+{
+    // todo:
+    for (auto &amf : m_amfCtx)
+        return amf.second; // return the first one
+    return nullptr;
+}
+
+NgapAmfContext *NgapTask::selectNewAmfForReAllocation(int ueId, int initiatedAmfId, int amfSetId)
+{
+    // TODO an arbitrary AMF is selected for now
+    return findAmfContext(initiatedAmfId);
+}
+
+} // namespace nr::gnb
diff --git a/src/gnb/ngap/task.hpp b/src/gnb/ngap/task.hpp
index 44e86ef17..df272ca8e 100644
--- a/src/gnb/ngap/task.hpp
+++ b/src/gnb/ngap/task.hpp
@@ -69,7 +69,6 @@ class NgapTask : public NtsTask
     NgapUeContext *findUeByRanId(long ranUeNgapId);
     NgapUeContext *findUeByAmfId(long amfUeNgapId);
     NgapUeContext *findUeByNgapIdPair(int amfCtxId, const NgapIdPair &idPair);
-    NgapAmfContext *selectNewAmfForReAllocation(int initiatedAmfId, int amfSetId);
     void deleteUeContext(int ueId);
     void deleteAmfContext(int amfId);
 
@@ -105,6 +104,10 @@ class NgapTask : public NtsTask
     void receiveInitialContextSetup(int amfId, ASN_NGAP_InitialContextSetupRequest *msg);
     void receiveContextRelease(int amfId, ASN_NGAP_UEContextReleaseCommand *msg);
     void receiveContextModification(int amfId, ASN_NGAP_UEContextModificationRequest *msg);
+
+    /* NAS Node Selection */
+    NgapAmfContext *selectAmf(int ueId);
+    NgapAmfContext *selectNewAmfForReAllocation(int ueId, int initiatedAmfId, int amfSetId);
 };
 
 } // namespace nr::gnb
\ No newline at end of file