From 469c3ad273cf558fe7f1ed4ab650b8d7f021eadf Mon Sep 17 00:00:00 2001
From: Pedro Alvarez <pedro.alvarez@softwareradiosystems.com>
Date: Tue, 25 Feb 2020 12:00:13 +0000
Subject: [PATCH] Removed more useless lookups.

---
 srsepc/hdr/hss/hss.h  |  8 ++++----
 srsepc/src/hss/hss.cc | 37 +++++++++++++++----------------------
 2 files changed, 19 insertions(+), 26 deletions(-)

diff --git a/srsepc/hdr/hss/hss.h b/srsepc/hdr/hss/hss.h
index b12a22534..794c3a282 100644
--- a/srsepc/hdr/hss/hss.h
+++ b/srsepc/hdr/hss/hss.h
@@ -92,17 +92,17 @@ private:
   void gen_rand(uint8_t rand_[16]);
   bool get_k_amf_opc_sqn(uint64_t imsi, uint8_t* k, uint8_t* amf, uint8_t* opc, uint8_t* sqn);
 
-  void gen_auth_info_answer_xor(hss_ue_ctx_t* ue_ctx, uint8_t* k_asme, uint8_t* autn, uint8_t* rand, uint8_t* xres);
   void gen_auth_info_answer_milenage(hss_ue_ctx_t* ue_ctx, uint8_t* k_asme, uint8_t* autn, uint8_t* rand, uint8_t* xres);
+  void gen_auth_info_answer_xor(hss_ue_ctx_t* ue_ctx, uint8_t* k_asme, uint8_t* autn, uint8_t* rand, uint8_t* xres);
 
-  bool resync_sqn_milenage(uint64_t imsi, uint8_t* auts);
-  bool resync_sqn_xor(uint64_t imsi, uint8_t* auts);
+  bool resync_sqn_milenage(hss_ue_ctx_t* ue_ctx, uint8_t* auts);
+  bool resync_sqn_xor(hss_ue_ctx_t* ue_ctx, uint8_t* auts);
 
   std::vector<std::string> split_string(const std::string& str, char delimiter);
   void                     get_uint_vec_from_hex_str(const std::string& key_str, uint8_t* key, uint len);
 
   void increment_ue_sqn(hss_ue_ctx_t* ue_ctx);
-  void increment_seq_after_resync(uint64_t imsi);
+  void increment_seq_after_resync(hss_ue_ctx_t* ue_ctx);
   void increment_sqn(uint8_t* sqn, uint8_t* next_sqn);
   void set_sqn(uint64_t imsi, uint8_t* sqn);
 
diff --git a/srsepc/src/hss/hss.cc b/srsepc/src/hss/hss.cc
index 9268c5d35..ab824c980 100644
--- a/srsepc/src/hss/hss.cc
+++ b/srsepc/src/hss/hss.cc
@@ -470,41 +470,39 @@ bool hss::resync_sqn(uint64_t imsi, uint8_t* auts)
   bool ret = false; //FIXME
   switch (ue_ctx->algo) {
     case HSS_ALGO_XOR:
-      ret = resync_sqn_xor(imsi, auts);
+      ret = resync_sqn_xor(ue_ctx, auts);
       break;
     case HSS_ALGO_MILENAGE:
-      ret = resync_sqn_milenage(imsi, auts);
+      ret = resync_sqn_milenage(ue_ctx, auts);
       break;
   }
 
-  increment_seq_after_resync(imsi);
+  increment_seq_after_resync(ue_ctx);
   return ret;
 }
 
-bool hss::resync_sqn_xor(uint64_t imsi, uint8_t* auts)
+bool hss::resync_sqn_xor(hss_ue_ctx_t* ue_ctx, uint8_t* auts)
 {
   m_hss_log->error("XOR SQN synchronization not supported yet\n");
   m_hss_log->console("XOR SQNs synchronization not supported yet\n");
   return false;
 }
 
-bool hss::resync_sqn_milenage(uint64_t imsi, uint8_t* auts)
+bool hss::resync_sqn_milenage(hss_ue_ctx_t* ue_ctx, uint8_t* auts)
 {
+  // Get K, AMF, OPC and SQN
+  uint8_t *k = ue_ctx->key;
+  uint8_t *amf = ue_ctx->amf;
+  uint8_t *opc = ue_ctx->opc;
+  uint8_t *sqn = ue_ctx->sqn;
+ 
+  // Temp variables
   uint8_t last_rand[16];
   uint8_t ak[6];
   uint8_t mac_s[8];
   uint8_t sqn_ms_xor_ak[6];
 
-  uint8_t k[16];
-  uint8_t amf[2];
-  uint8_t opc[16];
-  uint8_t sqn[6];
-
-  if (!get_k_amf_opc_sqn(imsi, k, amf, opc, sqn)) {
-    return false;
-  }
-
-  get_last_rand(imsi, last_rand);
+  get_last_rand(ue_ctx->imsi, last_rand);
 
   for (int i = 0; i < 6; i++) {
     sqn_ms_xor_ak[i] = auts[i];
@@ -542,7 +540,7 @@ bool hss::resync_sqn_milenage(uint64_t imsi, uint8_t* auts)
   srslte::security_milenage_f1_star(k, opc, last_rand, sqn_ms, amf, mac_s_tmp);
   m_hss_log->debug_hex(mac_s_tmp, 8, "MAC calc : ");
 
-  set_sqn(imsi, sqn_ms);
+  set_sqn(ue_ctx->imsi, sqn_ms);
   return true;
 }
 
@@ -583,14 +581,9 @@ void hss::increment_sqn(uint8_t* sqn, uint8_t* next_sqn)
   return;
 }
 
-void hss::increment_seq_after_resync(uint64_t imsi)
+void hss::increment_seq_after_resync(hss_ue_ctx_t* ue_ctx)
 {
   // This function only increment the SEQ part of the SQN for resynchronization purpose
-  hss_ue_ctx_t* ue_ctx = get_ue_ctx(imsi);
-  if (ue_ctx == nullptr) {
-    return;
-  }
-
   uint8_t* sqn = ue_ctx->sqn;
 
   uint64_t seq;