From a8eaf1cdd8b28827b00ae3db0d71ed5e4734cfee Mon Sep 17 00:00:00 2001
From: Bedran Karakoc <bedran.karakoc@srs.io>
Date: Thu, 20 Jan 2022 17:40:13 +0100
Subject: [PATCH] ue,nas_5g: Add get_k_amf() helper function and security
 context flag

---
 lib/include/srsran/interfaces/ue_nas_interfaces.h |  1 +
 srsue/hdr/stack/upper/nas_5g.h                    |  5 +++--
 srsue/src/stack/rrc_nr/test/ue_rrc_nr_test.cc     |  1 +
 srsue/src/stack/upper/nas_5g.cc                   | 15 +++++++++++++++
 4 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/lib/include/srsran/interfaces/ue_nas_interfaces.h b/lib/include/srsran/interfaces/ue_nas_interfaces.h
index 386270aa2..7cb81b666 100644
--- a/lib/include/srsran/interfaces/ue_nas_interfaces.h
+++ b/lib/include/srsran/interfaces/ue_nas_interfaces.h
@@ -59,6 +59,7 @@ class nas_5g_interface_rrc_nr
 {
 public:
   virtual int write_pdu(srsran::unique_byte_buffer_t pdu) = 0;
+  virtual int get_k_amf(srsran::as_key_t& k_amf)          = 0;
 };
 
 class nas_5g_interface_procedures
diff --git a/srsue/hdr/stack/upper/nas_5g.h b/srsue/hdr/stack/upper/nas_5g.h
index c0bf9e5fa..f317ca9ae 100644
--- a/srsue/hdr/stack/upper/nas_5g.h
+++ b/srsue/hdr/stack/upper/nas_5g.h
@@ -58,6 +58,7 @@ public:
 
   // Stack+RRC interface
   bool is_registered();
+  int  get_k_amf(srsran::as_key_t& k_amf);
 
   int write_pdu(srsran::unique_byte_buffer_t pdu);
 
@@ -79,8 +80,8 @@ private:
   usim_interface_nas*      usim   = nullptr;
   gw_interface_nas*        gw     = nullptr;
 
-  bool running = false;
-
+  bool                                             running             = false;
+  bool                                             has_sec_ctxt        = false;
   bool                                             initial_sec_command = false;
   srsran::nas_5g::mobile_identity_5gs_t::guti_5g_s guti_5g;
 
diff --git a/srsue/src/stack/rrc_nr/test/ue_rrc_nr_test.cc b/srsue/src/stack/rrc_nr/test/ue_rrc_nr_test.cc
index 29e03fbab..e4fa5e7a8 100644
--- a/srsue/src/stack/rrc_nr/test/ue_rrc_nr_test.cc
+++ b/srsue/src/stack/rrc_nr/test/ue_rrc_nr_test.cc
@@ -110,6 +110,7 @@ class dummy_eutra : public rrc_eutra_interface_rrc_nr
 class dummy_nas : public nas_5g_interface_rrc_nr
 {
   int write_pdu(srsran::unique_byte_buffer_t pdu) { return SRSRAN_SUCCESS; };
+  int get_k_amf(srsran::as_key_t& k_amf) { return SRSRAN_SUCCESS; };
 };
 
 class dummy_sim : public usim_interface_rrc_nr
diff --git a/srsue/src/stack/upper/nas_5g.cc b/srsue/src/stack/upper/nas_5g.cc
index ba2ac08d0..4dc468bd9 100644
--- a/srsue/src/stack/upper/nas_5g.cc
+++ b/srsue/src/stack/upper/nas_5g.cc
@@ -453,6 +453,7 @@ int nas_5g::send_security_mode_complete(const srsran::nas_5g::security_mode_comm
     pcap->write_nas(pdu.get()->msg, pdu.get()->N_bytes);
   }
 
+  has_sec_ctxt = true;
   logger.info("Sending Security Mode Complete");
   rrc_nr->write_sdu(std::move(pdu));
   ctxt_base.tx_count++;
@@ -798,6 +799,7 @@ int nas_5g::handle_registration_accept(registration_accept_t& registration_accep
 int nas_5g::handle_registration_reject(registration_reject_t& registration_reject)
 {
   logger.info("Handling Registration Reject");
+  has_sec_ctxt = false;
   ctxt_base.rx_count++;
   state.set_deregistered(mm5g_state_t::deregistered_substate_t::plmn_search);
 
@@ -888,6 +890,7 @@ int nas_5g::handle_authentication_request(authentication_request_t& authenticati
 int nas_5g::handle_authentication_reject(srsran::nas_5g::authentication_reject_t& authentication_reject)
 {
   logger.info("Handling Authentication Reject");
+  has_sec_ctxt = false;
   ctxt_base.rx_count++;
   state.set_deregistered(mm5g_state_t::deregistered_substate_t::plmn_search);
   return SRSRAN_SUCCESS;
@@ -911,6 +914,7 @@ int nas_5g::handle_service_accept(srsran::nas_5g::service_accept_t& service_acce
 int nas_5g::handle_service_reject(srsran::nas_5g::service_reject_t& service_reject)
 {
   logger.info("Handling Service Reject");
+  has_sec_ctxt = false;
   ctxt_base.rx_count++;
   return SRSRAN_SUCCESS;
 }
@@ -1106,6 +1110,17 @@ void nas_5g::get_metrics(nas_5g_metrics_t& metrics)
   metrics.state                   = state.get_state();
 }
 
+int nas_5g::get_k_amf(as_key_t& k_amf)
+{
+  if (not has_sec_ctxt) {
+    logger.error("K_amf requested before a valid NAS security context was established");
+    return SRSRAN_ERROR;
+  }
+
+  std::copy(std::begin(ctxt_5g.k_amf), std::end(ctxt_5g.k_amf), k_amf.begin());
+  return SRSRAN_SUCCESS;
+}
+
 /*******************************************************************************
  * Helpers
  ******************************************************************************/