gsup_server: Fix removal of session from set

Change-Id: Iec9bce7494e3849c9e94158516f471f97182725b
---
M src/gsup_server.erl
1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/gsup_server.erl b/src/gsup_server.erl
index fd2a8b5..8a070bf 100644
--- a/src/gsup_server.erl
+++ b/src/gsup_server.erl
@@ -198,7 +198,7 @@
 				}
 	end,
 	tx_gsup(Socket, Resp),
-	State1 = delete_gsups_ue(Imsi, State0),
+	State1 = delete_gsups_ue_by_imsi(Imsi, State0),
 	{noreply, State1};
 
 % Our GSUP CEAI implementation for "IKEv2 Information Delete Request"
@@ -244,7 +244,7 @@
 		},
 		tx_gsup(Socket, Resp),
 		epdg_ue_fsm:stop(UE#gsups_ue.pid),
-		State2 = delete_gsups_ue(Imsi, State1)
+		State2 = delete_gsups_ue(UE, State1)
 	end,
 	{noreply, State2};
 
@@ -331,7 +331,7 @@
 	lager:info("GSUP: Rx ~p~n", [GsupMsgRx]),
 	UE = find_gsups_ue_by_imsi(Imsi, State0),
 	case UE of
-	#gsups_ue{imsi = Imsi} -> State1 = delete_gsups_ue(Imsi, State0);
+	#gsups_ue{imsi = Imsi} -> State1 = delete_gsups_ue(UE, State0);
 	undefined -> State1 = State0
 	end,
 	{noreply, State1};
@@ -401,6 +401,13 @@
 		new_gsups_ue(Imsi, State)
 	end.
 
-delete_gsups_ue(Imsi, State) ->
-	SetRemoved = sets:del_element(Imsi, State#gsups_state.ues),
-	State#gsups_state{ues = SetRemoved}.
\ No newline at end of file
+delete_gsups_ue(UE, State) ->
+	SetRemoved = sets:del_element(UE, State#gsups_state.ues),
+	lager:debug("Removed UE ~p from ~p~n", [UE, SetRemoved]),
+	State#gsups_state{ues = SetRemoved}.
+
+delete_gsups_ue_by_imsi(Imsi, State) ->
+	case find_gsups_ue_by_imsi(Imsi, State) of
+	undefined -> State;
+	UE-> delete_gsups_ue(UE, State)
+	end.
\ No newline at end of file

