From e3d7ccd7e165eced845cb45e6ebc1ebe9f47078f Mon Sep 17 00:00:00 2001
From: Pedro Alvarez <pedro.alvarez@softwareradiosystems.com>
Date: Fri, 30 Nov 2018 11:36:21 +0000
Subject: [PATCH] Changed the way UL NAS count is obtained to generate K_enb in
 the UE. Incrementing UL NAS count on integrity protected dettach request.

---
 srsue/hdr/upper/nas.h  |  2 +-
 srsue/src/upper/nas.cc | 13 ++++++-------
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/srsue/hdr/upper/nas.h b/srsue/hdr/upper/nas.h
index 7d8cc8b01..685ac8b78 100644
--- a/srsue/hdr/upper/nas.h
+++ b/srsue/hdr/upper/nas.h
@@ -85,7 +85,6 @@ public:
   void paging(LIBLTE_RRC_S_TMSI_STRUCT *ue_identiy);
   void set_barring(barring_t barring);
   void write_pdu(uint32_t lcid, byte_buffer_t *pdu);
-  void set_k_enb_count();
   uint32_t get_k_enb_count();
   bool is_attached();
   bool get_k_asme(uint8_t *k_asme_, uint32_t n);
@@ -165,6 +164,7 @@ private:
   bool integrity_check(byte_buffer_t *pdu);
   void cipher_encrypt(byte_buffer_t *pdu);
   void cipher_decrypt(byte_buffer_t *pdu);
+  void set_k_enb_count(uint32_t count);
 
   bool check_cap_replay(LIBLTE_MME_UE_SECURITY_CAPABILITIES_STRUCT *caps);
 
diff --git a/srsue/src/upper/nas.cc b/srsue/src/upper/nas.cc
index cf2338b3f..8e46a988d 100644
--- a/srsue/src/upper/nas.cc
+++ b/srsue/src/upper/nas.cc
@@ -388,10 +388,10 @@ void nas::write_pdu(uint32_t lcid, byte_buffer_t *pdu) {
   }
 }
 
-void nas::set_k_enb_count() {
+void nas::set_k_enb_count(uint32_t count) {
   // UL count for RRC key derivation depends on UL count of the Authentication Request or Service Request.
   // This function should be called after sending these messages, for later derivation of the keys.
-  ctxt.k_enb_count = ctxt.tx_count;
+  ctxt.k_enb_count = count;
   return;
 }
 
@@ -889,7 +889,7 @@ void nas::parse_authentication_request(uint32_t lcid, byte_buffer_t *pdu, const
     nas_log->info("Network authentication successful\n");
     send_authentication_response(res, res_len, sec_hdr_type);
     nas_log->info_hex(ctxt.k_asme, 32, "Generated k_asme:\n");
-    set_k_enb_count();
+    set_k_enb_count(0);
     auth_request = true;
   } else if (auth_result == AUTH_SYNCH_FAILURE) {
     nas_log->error("Network authentication synchronization failure.\n");
@@ -1180,7 +1180,7 @@ void nas::gen_attach_request(byte_buffer_t *msg) {
   }
 
   if (have_ctxt) {
-    ctxt.tx_count++;
+    set_k_enb_count(ctxt.tx_count++);
   }
 }
 
@@ -1216,9 +1216,7 @@ void nas::gen_service_request(byte_buffer_t *msg) {
   if(pcap != NULL) {
     pcap->write_nas(msg->msg, msg->N_bytes);
   }
-
-  ctxt.tx_count++;
-  set_k_enb_count();
+  set_k_enb_count(ctxt.tx_count++);
 }
 
 void nas::gen_pdn_connectivity_request(LIBLTE_BYTE_MSG_STRUCT *msg) {
@@ -1325,6 +1323,7 @@ void nas::send_detach_request(bool switch_off)
                          &pdu->msg[5],
                          pdu->N_bytes - 5,
                          &pdu->msg[1]);
+      ctxt.tx_count++;
     } else {
       nas_log->error("Invalid PDU size %d\n", pdu->N_bytes);
     }