From f381e93237c411eb95c7bb0aa45393afbc3f6446 Mon Sep 17 00:00:00 2001
From: Sakthivel Velumani <mail@sakthi.me>
Date: Thu, 26 Oct 2023 14:30:36 -0400
Subject: [PATCH] Set NSSAI config in MAC Each active LCID at MAC is configured
 with NSSAI that is sent from AMF with PDU session setup request

---
 openair2/LAYER2/NR_MAC_gNB/mac_rrc_dl_handler.c | 17 +++++++++++++++++
 openair2/LAYER2/NR_MAC_gNB/nr_mac_gNB.h         |  1 +
 openair2/RRC/NR/rrc_gNB.c                       |  2 +-
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/openair2/LAYER2/NR_MAC_gNB/mac_rrc_dl_handler.c b/openair2/LAYER2/NR_MAC_gNB/mac_rrc_dl_handler.c
index ea0d08e152..90bce85c75 100644
--- a/openair2/LAYER2/NR_MAC_gNB/mac_rrc_dl_handler.c
+++ b/openair2/LAYER2/NR_MAC_gNB/mac_rrc_dl_handler.c
@@ -222,6 +222,17 @@ static NR_CellGroupConfig_t *clone_CellGroupConfig(const NR_CellGroupConfig_t *o
   return cloned;
 }
 
+static void set_nssaiConfig(const int drb_len, const f1ap_drb_to_be_setup_t *req_drbs, NR_UE_sched_ctrl_t *sched_ctrl)
+{
+  for (int i = 0; i < drb_len; i++) {
+    const f1ap_drb_to_be_setup_t *drb = &req_drbs[i];
+
+    long lcid = drb->drb_id + 3; /* LCID is DRB + 3 */
+    sched_ctrl->dl_lc_nssai[lcid] = drb->nssai;
+    LOG_I(NR_MAC, "Setting NSSAI sst: %d, sd: %d for DRB: %ld\n", drb->nssai.sst, drb->nssai.sd, drb->drb_id);
+  }
+}
+
 void ue_context_setup_request(const f1ap_ue_context_setup_t *req)
 {
   gNB_MAC_INST *mac = RC.nrmac[0];
@@ -288,6 +299,9 @@ void ue_context_setup_request(const f1ap_ue_context_setup_t *req)
   /* TODO: need to apply after UE context reconfiguration confirmed? */
   nr_mac_prepare_cellgroup_update(mac, UE, new_CellGroup);
 
+  /* Set NSSAI config in MAC for each active DRB */
+  set_nssaiConfig(req->drbs_to_be_setup_length, req->drbs_to_be_setup, &UE->UE_sched_ctrl);
+
   NR_SCHED_UNLOCK(&mac->sched_lock);
 
   /* some sanity checks, since we use the same type for request and response */
@@ -384,6 +398,9 @@ void ue_context_modification_request(const f1ap_ue_context_modif_req_t *req)
     resp.du_to_cu_rrc_information->cellGroupConfig_length = (enc_rval.encoded + 7) >> 3;
 
     nr_mac_prepare_cellgroup_update(mac, UE, new_CellGroup);
+
+    /* Set NSSAI config in MAC for each active DRB */
+    set_nssaiConfig(req->drbs_to_be_setup_length, req->drbs_to_be_setup, &UE->UE_sched_ctrl);
   } else {
     ASN_STRUCT_FREE(asn_DEF_NR_CellGroupConfig, new_CellGroup); // we actually don't need it
   }
diff --git a/openair2/LAYER2/NR_MAC_gNB/nr_mac_gNB.h b/openair2/LAYER2/NR_MAC_gNB/nr_mac_gNB.h
index 89f1dabc34..0060df1c10 100644
--- a/openair2/LAYER2/NR_MAC_gNB/nr_mac_gNB.h
+++ b/openair2/LAYER2/NR_MAC_gNB/nr_mac_gNB.h
@@ -639,6 +639,7 @@ typedef struct {
 
   /// sri, ul_ri and tpmi based on SRS
   nr_srs_feedback_t srs_feedback;
+  nssai_t dl_lc_nssai[NR_MAX_NUM_LCID];
 } NR_UE_sched_ctrl_t;
 
 typedef struct {
diff --git a/openair2/RRC/NR/rrc_gNB.c b/openair2/RRC/NR/rrc_gNB.c
index ec155d97bf..6c3de095e5 100644
--- a/openair2/RRC/NR/rrc_gNB.c
+++ b/openair2/RRC/NR/rrc_gNB.c
@@ -2346,7 +2346,7 @@ void rrc_gNB_process_e1_bearer_context_setup_resp(e1ap_bearer_setup_resp_t *resp
   /* Instruction towards the DU for DRB configuration and tunnel creation */
   int nb_drb = resp->pduSession[0].numDRBSetup;
   f1ap_drb_to_be_setup_t drbs[nb_drb];
-  rrc_pdu_session_param_t *RRC_pduSession = find_pduSession(UE, e1ap_resp->pduSession[0].id, false);
+  rrc_pdu_session_param_t *RRC_pduSession = find_pduSession(UE, resp->pduSession[0].id, false);
   DevAssert(RRC_pduSession);
   for (int i = 0; i < nb_drb; i++) {
     drbs[i].drb_id = resp->pduSession[0].DRBnGRanList[i].id;