S-TMSI handling in gNodeB RRC connection setup

src/gnb/rrc/connection.cpp

@@ -49,28 +49,26 @@ void GnbRrcTask::receiveRrcSetupRequest(int ueId, const ASN_RRC_RRCSetupRequest
         return;
     }
 
-    if (msg.rrcSetupRequest.ue_Identity.present == ASN_RRC_InitialUE_Identity_PR_ng_5G_S_TMSI_Part1)
+    if (msg.rrcSetupRequest.ue_Identity.present == ASN_RRC_InitialUE_Identity_PR_NOTHING)
     {
-        m_logger->err("RRC Setup Request with TMSI not implemented yet");
+        m_logger->err("Bad constructed RRC message ignored");
         return;
     }
 
-    if (msg.rrcSetupRequest.ue_Identity.present != ASN_RRC_InitialUE_Identity_PR_randomValue)
+    ue = createUe(ueId);
+
+    if (msg.rrcSetupRequest.ue_Identity.present == ASN_RRC_InitialUE_Identity_PR_ng_5G_S_TMSI_Part1)
     {
-        m_logger->err("Bad constructed RRC message ignored");
-        return;
+        ue->initialId = asn::GetBitStringLong<39>(msg.rrcSetupRequest.ue_Identity.choice.ng_5G_S_TMSI_Part1);
+        ue->isInitialIdSTmsi = true;
     }
-
-    int64_t initialRandomId = asn::GetBitStringLong<39>(msg.rrcSetupRequest.ue_Identity.choice.randomValue);
-    if (tryFindByInitialRandomId(initialRandomId) != nullptr)
+    else
     {
-        m_logger->err("Initial random ID conflict [%ld], discarding RRC Setup Request", initialRandomId);
-        return;
+        ue->initialId = asn::GetBitStringLong<39>(msg.rrcSetupRequest.ue_Identity.choice.randomValue);
+        ue->isInitialIdSTmsi = false;
     }
 
-    ue = createUe(ueId);
-    ue->initialRandomId = initialRandomId;
-    ue->establishmentCause = msg.rrcSetupRequest.establishmentCause;
+    ue->establishmentCause = static_cast<long>(msg.rrcSetupRequest.establishmentCause);
 
     // Prepare RRC Setup
     auto *pdu = asn::New<ASN_RRC_DL_CCCH_Message>();
@@ -100,11 +98,34 @@ void GnbRrcTask::receiveRrcSetupComplete(int ueId, const ASN_RRC_RRCSetupComplet
 
     auto setupComplete = msg.criticalExtensions.choice.rrcSetupComplete;
 
+    if (msg.criticalExtensions.choice.rrcSetupComplete)
+    {
+        // Handle received 5G S-TMSI if any
+        if (msg.criticalExtensions.choice.rrcSetupComplete->ng_5G_S_TMSI_Value)
+        {
+            ue->sTmsi = std::nullopt;
+
+            auto &sTmsiValue = msg.criticalExtensions.choice.rrcSetupComplete->ng_5G_S_TMSI_Value;
+            if (sTmsiValue->present == ASN_RRC_RRCSetupComplete_IEs__ng_5G_S_TMSI_Value_PR_ng_5G_S_TMSI)
+            {
+                ue->sTmsi = GutiMobileIdentity::FromSTmsi(asn::GetBitStringLong<48>(sTmsiValue->choice.ng_5G_S_TMSI));
+            }
+            else if (sTmsiValue->present == ASN_RRC_RRCSetupComplete_IEs__ng_5G_S_TMSI_Value_PR_ng_5G_S_TMSI_Part2)
+            {
+                if (ue->isInitialIdSTmsi)
+                {
+                    int64_t part2 = asn::GetBitStringLong<9>(sTmsiValue->choice.ng_5G_S_TMSI_Part2);
+                    ue->sTmsi = GutiMobileIdentity::FromSTmsi((part2 << 39) | (ue->initialId));
+                }
+            }
+        }
+    }
+
     auto *w = new NmGnbRrcToNgap(NmGnbRrcToNgap::INITIAL_NAS_DELIVERY);
     w->ueId = ueId;
     w->pdu = asn::GetOctetString(setupComplete->dedicatedNAS_Message);
     w->rrcEstablishmentCause = ue->establishmentCause;
-    w->sTmsi = std::nullopt; // TODO
+    w->sTmsi = ue->sTmsi;
 
     m_base->ngapTask->push(w);
 }

src/gnb/rrc/management.cpp

@@ -11,17 +11,6 @@
 namespace nr::gnb
 {
 
-RrcUeContext *GnbRrcTask::tryFindByInitialRandomId(int64_t id)
-{
-    if (id == -1)
-        return nullptr;
-    // TODO: Optimize
-    for (auto &item : m_ueCtx)
-        if (item.second->initialRandomId == id)
-            return item.second;
-    return nullptr;
-}
-
 int GnbRrcTask::getNextTid()
 {
     m_tidCounter++;

src/gnb/rrc/task.hpp

@@ -66,7 +66,6 @@ class GnbRrcTask : public NtsTask
 
   private:
     /* Management */
-    RrcUeContext *tryFindByInitialRandomId(int64_t id);
     int getNextTid();
 
     /* Handlers */

src/gnb/types.hpp

@@ -145,8 +145,10 @@ struct RrcUeContext
 {
     const int ueId{};
 
-    int64_t initialRandomId = -1;
+    int64_t initialId = -1; // 39-bit value, or -1
+    bool isInitialIdSTmsi{}; // TMSI-part-1 or a random value
     long establishmentCause{};
+    std::optional<GutiMobileIdentity> sTmsi{};
 
     explicit RrcUeContext(const int ueId) : ueId(ueId)
     {

src/utils/common_types.cpp

@@ -29,6 +29,15 @@ Supi Supi::Parse(const std::string &supi)
     throw std::runtime_error("invalid SUPI value");
 }
 
+GutiMobileIdentity GutiMobileIdentity::FromSTmsi(int64_t sTmsi)
+{
+    GutiMobileIdentity res;
+    res.tmsi = octet4{static_cast<uint32_t>(sTmsi & 0xFFFFFFLL)};
+    res.amfPointer = static_cast<int>(((sTmsi >> 32LL) & 0b111111LL));
+    res.amfSetId = static_cast<int>(((sTmsi >> 38LL) & 0b1111111111LL));
+    return res;
+}
+
 Json ToJson(const Supi &v)
 {
     return v.type + "-" + v.value;

src/utils/common_types.hpp

@@ -90,6 +90,8 @@ struct GutiMobileIdentity
         : plmn(plmn), amfRegionId(amfRegionId), amfSetId(amfSetId), amfPointer(amfPointer), tmsi(tmsi)
     {
     }
+
+    static GutiMobileIdentity FromSTmsi(int64_t sTmsi);
 };
 
 struct ImsiMobileIdentity