NAS Count handling in case of an RRC establishment failure

src/ue/nas/mm/radio.cpp

@@ -216,13 +216,29 @@ void NasMm::handleRrcConnectionRelease()
 
 void NasMm::handleRrcEstablishmentFailure()
 {
-    // TODO
-
     m_logger->err("RRC Establishment failure");
 
+    /* Handle NAS count */
+    {
+        // "After each new or retransmitted outbound SECURITY PROTECTED 5GS NAS MESSAGE message, the sender shall
+        // increase the NAS COUNT number by one, except for the initial NAS messages if the lower layers indicated the
+        // failure to establish the RRC connection"
+        bool hasNsCtx = m_usim->m_currentNsCtx &&
+                        (m_usim->m_currentNsCtx->integrity != nas::ETypeOfIntegrityProtectionAlgorithm::IA0 ||
+                         m_usim->m_currentNsCtx->ciphering != nas::ETypeOfCipheringAlgorithm::EA0);
+        if (hasNsCtx && m_usim->m_currentNsCtx->uplinkCount.sqn != 0 &&
+            m_usim->m_currentNsCtx->uplinkCount.overflow.operator int() != 0)
+        {
+            m_usim->m_currentNsCtx->rollbackCountOnEncrypt();
+        }
+    }
+
+    /* Handle MM state changes */
+    {
         if (m_mmState == EMmState::MM_REGISTERED_INITIATED)
         {
-        switchMmState(m_rmState == ERmState::RM_REGISTERED ? EMmSubState::MM_REGISTERED_ATTEMPTING_REGISTRATION_UPDATE
+            switchMmState(m_rmState == ERmState::RM_REGISTERED
+                              ? EMmSubState::MM_REGISTERED_ATTEMPTING_REGISTRATION_UPDATE
                               : EMmSubState::MM_DEREGISTERED_INITIAL_REGISTRATION_NEEDED);
         }
         else if (m_mmState == EMmState::MM_SERVICE_REQUEST_INITIATED)
@@ -232,6 +248,7 @@ void NasMm::handleRrcEstablishmentFailure()
         else if (m_mmState == EMmState::MM_DEREGISTERED_INITIATED)
         {
         }
+    }
 }
 
 void NasMm::handleRadioLinkFailure()

src/ue/types.hpp

@@ -446,6 +446,23 @@ struct NasSecurityContext
             uplinkCount.overflow = octet2(((int)uplinkCount.overflow + 1) & 0xFFFF);
     }
 
+    void rollbackCountOnEncrypt()
+    {
+        if (uplinkCount.sqn == 0)
+        {
+            uplinkCount.sqn = 0xFF;
+
+            if ((int)uplinkCount.overflow == 0)
+                uplinkCount.overflow = octet2{0xFFFF};
+            else
+                uplinkCount.overflow = octet2{(int)uplinkCount.overflow - 1};
+        }
+        else
+        {
+            uplinkCount.sqn = static_cast<uint8_t>(((int)uplinkCount.sqn - 1) & 0xFF);
+        }
+    }
+
     [[nodiscard]] NasSecurityContext deepCopy() const
     {
         NasSecurityContext ctx;