From ecb681cfd25e40505e770d56d1598445c10b58a3 Mon Sep 17 00:00:00 2001
From: Spencer Sevilla <spencer.builds.networks@gmail.com>
Date: Thu, 18 Jan 2024 14:11:09 -0800
Subject: [PATCH] need to be more defensive with sending downlinkdataack

---
 src/mme/mme-gtp-path.c    | 5 +++++
 src/mme/mme-s11-handler.c | 6 ++++--
 src/mme/s1ap-handler.c    | 9 +++++++++
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/mme/mme-gtp-path.c b/src/mme/mme-gtp-path.c
index dec42e66b6..35f1f3e120 100644
--- a/src/mme/mme-gtp-path.c
+++ b/src/mme/mme-gtp-path.c
@@ -603,6 +603,11 @@ int mme_gtp_send_downlink_data_notification_ack(
         ogs_warn("GTP transaction(NOTIFY) has already been removed");
         return OGS_OK;
     }
+    if (xact->step != 1) {
+        ogs_error("Trying to re-send downlink data notification ack?!?");
+        return OGS_OK;        
+    }
+
     mme_ue = bearer->mme_ue;
     ogs_assert(mme_ue);
     sgw_ue = mme_ue->sgw_ue;
diff --git a/src/mme/mme-s11-handler.c b/src/mme/mme-s11-handler.c
index 4982d80c17..f35f405623 100644
--- a/src/mme/mme-s11-handler.c
+++ b/src/mme/mme-s11-handler.c
@@ -1574,8 +1574,10 @@ void mme_s11_handle_downlink_data_notification(
             enb_ue_t *enb_ue = enb_ue_cycle(mme_ue->enb_ue);
             ogs_assert(enb_ue);
 
-            MME_STORE_PAGING_INFO(mme_ue,
-                MME_PAGING_TYPE_DOWNLINK_DATA_NOTIFICATION, bearer);            
+            mme_ue->paging.data = bearer;
+            // we are not calling MME_STORE_PAGING_INFO because we
+            // have not yet started paging, but we still need to keep
+            // a pointer to the relevant bearer to build the ACK later
 
             r = s1ap_send_ue_context_release_command(enb_ue,
                     S1AP_Cause_PR_nas, S1AP_CauseNas_normal_release,
diff --git a/src/mme/s1ap-handler.c b/src/mme/s1ap-handler.c
index 949580f800..4d69a44224 100644
--- a/src/mme/s1ap-handler.c
+++ b/src/mme/s1ap-handler.c
@@ -1663,6 +1663,15 @@ void s1ap_handle_ue_context_release_action(enb_ue_t *enb_ue)
         }
         enb_ue_unlink(mme_ue);
 
+        mme_bearer_t *bearer = mme_bearer_cycle(mme_ue->paging.data);
+        if (!bearer) {
+            ogs_error("Bearer is outdated, cannot page");
+            MME_CLEAR_PAGING_INFO(mme_ue);
+            return;
+        }
+
+        MME_STORE_PAGING_INFO(mme_ue,
+            MME_PAGING_TYPE_DOWNLINK_DATA_NOTIFICATION, bearer);
         r = s1ap_send_paging(mme_ue, S1AP_CNDomain_ps);
         ogs_expect(r == OGS_OK);
         ogs_assert(r != OGS_ERROR);