From 610b9980388e88b5c53193d712db07c21002491b Mon Sep 17 00:00:00 2001
From: Sukchan Lee <acetcom@gmail.com>
Date: Sun, 29 Jan 2023 09:09:52 +0900
Subject: [PATCH] [AMX] Fixed a crash due to deregistration (#2021)

Fixed an issue where AMF would crash
if an implicit deregistration occurred twice.
---
 src/amf/amf-sm.c | 35 ++++++++++++++++++---------------
 src/amf/amf-sm.h |  1 +
 src/amf/gmm-sm.c | 50 +++++++++++++++++++++++++++++++++++++++++++++---
 3 files changed, 68 insertions(+), 18 deletions(-)

diff --git a/src/amf/amf-sm.c b/src/amf/amf-sm.c
index 95a4c4c4d..265115674 100644
--- a/src/amf/amf-sm.c
+++ b/src/amf/amf-sm.c
@@ -378,18 +378,18 @@ void amf_state_operational(ogs_fsm_t *s, amf_event_t *e)
             ogs_sbi_xact_remove(sbi_xact);
 
             amf_ue = amf_ue_cycle(amf_ue);
-
-            if (amf_ue) {
-                ogs_assert(OGS_FSM_STATE(&amf_ue->sm));
-
-                e->amf_ue = amf_ue;
-                e->h.sbi.message = &sbi_message;;
-                e->h.sbi.state = state;
-
-                ogs_fsm_dispatch(&amf_ue->sm, e);
-            } else {
+            if (!amf_ue) {
                 ogs_error("UE(amf_ue) Context has already been removed");
+                break;
             }
+
+            ogs_assert(OGS_FSM_STATE(&amf_ue->sm));
+
+            e->amf_ue = amf_ue;
+            e->h.sbi.message = &sbi_message;;
+            e->h.sbi.state = state;
+
+            ogs_fsm_dispatch(&amf_ue->sm, e);
             break;
 
         CASE(OGS_SBI_SERVICE_NAME_NSMF_PDUSESSION)
@@ -807,7 +807,7 @@ void amf_state_operational(ogs_fsm_t *s, amf_event_t *e)
         if (ogs_nas_5gmm_decode(&nas_message, pkbuf) != OGS_OK) {
             ogs_error("ogs_nas_5gmm_decode() failed");
             ogs_pkbuf_free(pkbuf);
-            return;
+            break;
         }
 
         amf_ue = ran_ue->amf_ue;
@@ -824,7 +824,7 @@ void amf_state_operational(ogs_fsm_t *s, amf_event_t *e)
                     ogs_expect(r == OGS_OK);
                     ogs_assert(r != OGS_ERROR);
                     ogs_pkbuf_free(pkbuf);
-                    return;
+                    break;
                 }
             } else {
                 /* Here, if the AMF_UE Context is found,
@@ -844,7 +844,7 @@ void amf_state_operational(ogs_fsm_t *s, amf_event_t *e)
                         ogs_error("[%s] nas_security_decode() failed",
                                 amf_ue->suci);
                         ogs_pkbuf_free(pkbuf);
-                        return;
+                        break;
                     }
                 }
             }
@@ -909,9 +909,14 @@ void amf_state_operational(ogs_fsm_t *s, amf_event_t *e)
 
         ogs_pkbuf_free(pkbuf);
         break;
+
     case AMF_EVENT_5GMM_TIMER:
-        amf_ue = e->amf_ue;
-        ogs_assert(amf_ue);
+        amf_ue = amf_ue_cycle(e->amf_ue);
+        if (!amf_ue) {
+            ogs_error("UE(amf_ue) Context has already been removed");
+            break;
+        }
+
         ogs_assert(OGS_FSM_STATE(&amf_ue->sm));
 
         ogs_fsm_dispatch(&amf_ue->sm, e);
diff --git a/src/amf/amf-sm.h b/src/amf/amf-sm.h
index 81f214aff..5cd19c0fe 100644
--- a/src/amf/amf-sm.h
+++ b/src/amf/amf-sm.h
@@ -42,6 +42,7 @@ void gmm_state_authentication(ogs_fsm_t *s, amf_event_t *e);
 void gmm_state_security_mode(ogs_fsm_t *s, amf_event_t *e);
 void gmm_state_initial_context_setup(ogs_fsm_t *s, amf_event_t *e);
 void gmm_state_registered(ogs_fsm_t *s, amf_event_t *e);
+void gmm_state_ue_context_will_remove(ogs_fsm_t *s, amf_event_t *e);
 void gmm_state_exception(ogs_fsm_t *s, amf_event_t *e);
 
 #define amf_sm_debug(__pe) \
diff --git a/src/amf/gmm-sm.c b/src/amf/gmm-sm.c
index 19f30a3b9..56a2b3ad8 100644
--- a/src/amf/gmm-sm.c
+++ b/src/amf/gmm-sm.c
@@ -60,8 +60,11 @@ void gmm_state_de_registered(ogs_fsm_t *s, amf_event_t *e)
 
     int r, state = 0;
 
+    ogs_assert(s);
     ogs_assert(e);
 
+    amf_sm_debug(e);
+
     if (e->sess) {
         sess = e->sess;
         amf_ue = sess->amf_ue;
@@ -289,13 +292,17 @@ void gmm_state_de_registered(ogs_fsm_t *s, amf_event_t *e)
                     break;
 
                 CASE(OGS_SBI_HTTP_METHOD_DELETE)
-                    if (state != AMF_NETWORK_INITIATED_DE_REGISTERED) {
+                    if (state == AMF_NETWORK_INITIATED_DE_REGISTERED) {
+                        ogs_warn("[%s] AMF-UE Context Removed", amf_ue->supi);
+                        OGS_FSM_TRAN(&amf_ue->sm,
+                                &gmm_state_ue_context_will_remove);
+                    } else {
                         r = nas_5gs_send_de_registration_accept(amf_ue);
                         ogs_expect(r == OGS_OK);
                         ogs_assert(r != OGS_ERROR);
-                    }
 
-                    PCF_AM_POLICY_CLEAR(amf_ue);
+                        PCF_AM_POLICY_CLEAR(amf_ue);
+                    }
                     break;
 
                 DEFAULT
@@ -331,8 +338,11 @@ void gmm_state_registered(ogs_fsm_t *s, amf_event_t *e)
 
     ogs_sbi_message_t *sbi_message = NULL;
 
+    ogs_assert(s);
     ogs_assert(e);
 
+    amf_sm_debug(e);
+
     if (e->sess) {
         sess = e->sess;
         amf_ue = sess->amf_ue;
@@ -1578,6 +1588,38 @@ void gmm_state_initial_context_setup(ogs_fsm_t *s, amf_event_t *e)
     }
 }
 
+void gmm_state_ue_context_will_remove(ogs_fsm_t *s, amf_event_t *e)
+{
+    amf_ue_t *amf_ue = NULL;
+    amf_sess_t *sess = NULL;
+
+    ogs_assert(s);
+    ogs_assert(e);
+
+    amf_sm_debug(e);
+
+    if (e->sess) {
+        sess = e->sess;
+        amf_ue = sess->amf_ue;
+        ogs_assert(amf_ue);
+    } else {
+        amf_ue = e->amf_ue;
+        ogs_assert(amf_ue);
+    }
+
+    switch (e->h.id) {
+    case OGS_FSM_ENTRY_SIG:
+        amf_ue_remove(amf_ue);
+        break;
+
+    case OGS_FSM_EXIT_SIG:
+        break;
+
+    default:
+        ogs_error("Unknown event[%s]", amf_event_get_name(e));
+    }
+}
+
 void gmm_state_exception(ogs_fsm_t *s, amf_event_t *e)
 {
     int xact_count = 0, r;
@@ -1589,7 +1631,9 @@ void gmm_state_exception(ogs_fsm_t *s, amf_event_t *e)
     ogs_nas_5gs_message_t *nas_message = NULL;
     ogs_nas_security_header_type_t h;
 
+    ogs_assert(s);
     ogs_assert(e);
+
     amf_sm_debug(e);
 
     if (e->sess) {