diff --git a/srsepc/hdr/mme/s1ap.h b/srsepc/hdr/mme/s1ap.h
index b2c741d26..d417d9442 100644
--- a/srsepc/hdr/mme/s1ap.h
+++ b/srsepc/hdr/mme/s1ap.h
@@ -80,8 +80,8 @@ public:
   enb_ctx_t* find_enb_ctx(uint16_t enb_id);
   void add_new_enb_ctx(const enb_ctx_t &enb_ctx, const struct sctp_sndrcvinfo* enb_sri);
 
-  bool add_ue_ctx_to_imsi_map(const ue_ctx_t *ue_ctx);
-  bool add_ue_ctx_to_mme_ue_s1ap_id_map(const ue_ctx_t *ue_ctx);
+  bool add_ue_ctx_to_imsi_map(ue_ctx_t *ue_ctx);
+  bool add_ue_ctx_to_mme_ue_s1ap_id_map(ue_ctx_t *ue_ctx);
 
   ue_ctx_t* find_ue_ctx_from_imsi(uint64_t imsi);
   ue_ctx_t* find_ue_ctx_from_mme_ue_s1ap_id(uint32_t mme_ue_s1ap_id);
diff --git a/srsepc/src/mme/s1ap.cc b/srsepc/src/mme/s1ap.cc
index a541f00e3..4aeb17011 100644
--- a/srsepc/src/mme/s1ap.cc
+++ b/srsepc/src/mme/s1ap.cc
@@ -384,7 +384,7 @@ s1ap::add_new_ue_ctx(const ue_ctx_t &ue_ctx)
 }
 */
 bool
-s1ap::add_ue_ctx_to_imsi_map(const ue_ctx_t *ue_ctx)
+s1ap::add_ue_ctx_to_imsi_map(ue_ctx_t *ue_ctx)
 {
   std::map<uint64_t, ue_ctx_t*>::iterator ctx_it = m_imsi_to_ue_ctx.find(ue_ctx->emm_ctx.imsi);
   if(ctx_it != m_imsi_to_ue_ctx.end())
@@ -401,11 +401,13 @@ s1ap::add_ue_ctx_to_imsi_map(const ue_ctx_t *ue_ctx)
       return false;
     }
   }
+  m_imsi_to_ue_ctx.insert(std::pair<uint64_t,ue_ctx_t*>(ue_ctx->emm_ctx.imsi, ue_ctx));
+  m_s1ap_log->debug("Saved UE context corresponding to IMSI %015lu\n",ue_ctx->emm_ctx.imsi);
   return true;
 }
 
 bool
-s1ap::add_ue_ctx_to_mme_ue_s1ap_id_map(const ue_ctx_t *ue_ctx)
+s1ap::add_ue_ctx_to_mme_ue_s1ap_id_map(ue_ctx_t *ue_ctx)
 {
   if(ue_ctx->ecm_ctx.mme_ue_s1ap_id == 0)
   {
@@ -427,6 +429,8 @@ s1ap::add_ue_ctx_to_mme_ue_s1ap_id_map(const ue_ctx_t *ue_ctx)
       return false;
     }
   }
+  m_mme_ue_s1ap_id_to_ue_ctx.insert(std::pair<uint32_t,ue_ctx_t*>(ue_ctx->ecm_ctx.mme_ue_s1ap_id, ue_ctx));
+  m_s1ap_log->debug("Saved UE context corresponding to MME UE S1AP Id %d\n",ue_ctx->ecm_ctx.mme_ue_s1ap_id);
   return true;
 }
 
diff --git a/srsepc/src/mme/s1ap_nas_transport.cc b/srsepc/src/mme/s1ap_nas_transport.cc
index 336d84143..135e99ae2 100644
--- a/srsepc/src/mme/s1ap_nas_transport.cc
+++ b/srsepc/src/mme/s1ap_nas_transport.cc
@@ -136,7 +136,6 @@ s1ap_nas_transport::handle_uplink_nas_transport(LIBLTE_S1AP_MESSAGE_UPLINKNASTRA
   uint8_t pd, msg_type, sec_hdr_type;
   uint32_t enb_ue_s1ap_id = ul_xport->eNB_UE_S1AP_ID.ENB_UE_S1AP_ID;
   uint32_t mme_ue_s1ap_id = ul_xport->MME_UE_S1AP_ID.MME_UE_S1AP_ID;
-  ue_emm_ctx_t *ue_emm_ctx = NULL;
   bool mac_valid = false;
 
   //Get UE ECM context
@@ -231,7 +230,7 @@ s1ap_nas_transport::handle_uplink_nas_transport(LIBLTE_S1AP_MESSAGE_UPLINKNASTRA
   {
     //Integrity protected NAS message, possibly chiphered.
     emm_ctx->security_ctxt.ul_nas_count++;
-    mac_valid = integrity_check(ue_emm_ctx,nas_msg);
+    mac_valid = integrity_check(emm_ctx,nas_msg);
     if(!mac_valid){
       m_s1ap_log->warning("Invalid MAC in NAS message type 0x%x.\n", msg_type);
       m_pool->deallocate(nas_msg);
@@ -273,14 +272,8 @@ s1ap_nas_transport::handle_uplink_nas_transport(LIBLTE_S1AP_MESSAGE_UPLINKNASTRA
 
   if(*reply_flag == true)
   {
-    if(ue_emm_ctx != NULL){
-      m_s1ap_log->console("DL NAS: Sent Downlink NAs Message. DL NAS Count=%d, UL NAS count=%d\n",emm_ctx->security_ctxt.dl_nas_count,emm_ctx->security_ctxt.ul_nas_count );
-      m_s1ap_log->info("DL NAS: Sent Downlink NAS message. DL NAS Count=%d, UL NAS count=%d\n",emm_ctx->security_ctxt.dl_nas_count, emm_ctx->security_ctxt.ul_nas_count);
-    }
-    else{
-      m_s1ap_log->console("DL NAS: Sent Downlink NAS Message\n");
-      m_s1ap_log->console("DL NAS: Sent Downlink NAS Message\n");
-    }
+    m_s1ap_log->console("DL NAS: Sent Downlink NAs Message. DL NAS Count=%d, UL NAS count=%d\n",emm_ctx->security_ctxt.dl_nas_count,emm_ctx->security_ctxt.ul_nas_count );
+    m_s1ap_log->info("DL NAS: Sent Downlink NAS message. DL NAS Count=%d, UL NAS count=%d\n",emm_ctx->security_ctxt.dl_nas_count, emm_ctx->security_ctxt.ul_nas_count);
   }
   m_pool->deallocate(nas_msg);
   return true;