From 68d4a909adbf59132d1d6d84966df4bc86b8c18a Mon Sep 17 00:00:00 2001
From: Andreas Schultz <andreas.schultz@travelping.com>
Date: Wed, 16 Dec 2020 12:32:58 +0100
Subject: [PATCH] replace context key tupels with records

Instead of having to count fields in tuples, records allow us to
specify matches with the more readable record syntax.
They also give us `tagged keys`, that means we can deduce the type
of the key from the record tag.

This was originally part of the stateless work, but will also be
helpful for the cluster registry.
---
 include/ergw.hrl          |  4 +++
 src/ergw.erl              |  3 +-
 src/ergw_api.erl          |  6 ++--
 src/ergw_context.erl      | 28 +++++++--------
 src/ergw_gtp_socket.erl   |  5 +--
 src/ergw_pfcp.erl         |  2 +-
 src/ergw_pfcp_context.erl |  2 +-
 src/ergw_sx_node.erl      |  2 +-
 src/gtp_context.erl       | 30 ++++++++--------
 src/gtp_context_reg.erl   | 14 --------
 src/gtp_v1_c.erl          | 16 +++------
 src/gtp_v2_c.erl          | 16 +++------
 test/ergw_test_lib.erl    |  4 ++-
 test/ergw_test_sx_up.erl  |  3 +-
 test/ggsn_SUITE.erl       | 37 +++++++++++++-------
 test/ggsn_proxy_SUITE.erl | 73 +++++++++++++++++++++++++++------------
 test/pgw_SUITE.erl        | 53 +++++++++++++++++++---------
 test/pgw_proxy_SUITE.erl  | 73 ++++++++++++++++++++++++++++-----------
 test/saegw_s11_SUITE.erl  | 37 ++++++++++++++------
 test/tdf_SUITE.erl        |  7 ++--
 20 files changed, 250 insertions(+), 165 deletions(-)

diff --git a/include/ergw.hrl b/include/ergw.hrl
index 407c434b..060b18d2 100644
--- a/include/ergw.hrl
+++ b/include/ergw.hrl
@@ -172,3 +172,7 @@
 	  ip                        :: {inet:ip4_address(),1..32}|
 				       {inet:ip6_address(),1..128}
 	}).
+
+-record(seid_key, {seid}).
+-record(context_key, {socket, id}).
+-record(socket_teid_key, {name, type, teid}).
diff --git a/src/ergw.erl b/src/ergw.erl
index f84cda2f..69283c50 100644
--- a/src/ergw.erl
+++ b/src/ergw.erl
@@ -27,6 +27,7 @@
 	 terminate/2, code_change/3]).
 
 -include_lib("kernel/include/logger.hrl").
+-include("include/ergw.hrl").
 
 -define(SERVER, ?MODULE).
 -record(state, {tid :: ets:tid()}).
@@ -208,7 +209,7 @@ i(memory, path) ->
 
 i(memory, context) ->
     MemUsage =
-	lists:foldl(fun({{seid, _}, {_, Pid}}, Mem) ->
+	lists:foldl(fun({#seid_key{}, {_, Pid}}, Mem) ->
 			    {memory, M} = erlang:process_info(Pid, memory),
 			    Mem + M;
 		       (_, Mem) ->
diff --git a/src/ergw_api.erl b/src/ergw_api.erl
index a411cd35..1844fd91 100644
--- a/src/ergw_api.erl
+++ b/src/ergw_api.erl
@@ -14,6 +14,8 @@
 
 -ignore_xref([peer/1, tunnel/1, memory/1]).
 
+-include("include/ergw.hrl").
+
 %%%===================================================================
 %%% API
 %%%===================================================================
@@ -38,8 +40,8 @@ tunnel(Socket) when is_atom(Socket) ->
     lists:foldl(fun collext_path_contexts/2, [], gtp_path_reg:all(Socket)).
 
 contexts(all) ->
-    lists:usort([Pid || {{_Socket, {teid, 'gtp-c', _TEID}}, {_, Pid}}
-				       <- gtp_context_reg:all(), is_pid(Pid)]).
+    lists:usort([Pid || {#socket_teid_key{type = 'gtp-c'}, {_, Pid}}
+			    <- gtp_context_reg:all(), is_pid(Pid)]).
 
 delete_contexts(all) ->
     lists:foreach(fun(Context) ->
diff --git a/src/ergw_context.erl b/src/ergw_context.erl
index 69fcab09..f2b232ea 100644
--- a/src/ergw_context.erl
+++ b/src/ergw_context.erl
@@ -39,7 +39,7 @@
 %%% -----------------------------------------------------------------
 
 sx_report(#pfcp{type = session_report_request, seid = SEID} = Report) ->
-    apply2context({seid, SEID}, sx_report, [Report]).
+    apply2context(#seid_key{seid = SEID}, sx_report, [Report]).
 
 %% port_message/2
 port_message(Request, Msg) ->
@@ -47,10 +47,15 @@ port_message(Request, Msg) ->
     ok.
 
 %% port_message/3
-port_message(Keys, #request{socket = Socket} = Request, Msg)
-  when is_list(Keys) ->
-    Contexts = gtp_context_reg:match_keys(Socket, Keys),
-    port_message_ctx(Contexts, Request, Msg).
+port_message(Id, #request{socket = Socket} = Request, Msg) ->
+    Key = gtp_context:context_key(Socket, Id),
+    case gtp_context_reg:select(Key) of
+	[{Handler, Server}] when is_atom(Handler), is_pid(Server) ->
+	    Handler:port_message(Server, Request, Msg, false);
+	_Other ->
+	    ?LOG(debug, "unable to find context ~p", [Key]),
+	    throw({error, not_found})
+    end.
 
 %% port_message/4
 port_message(Key, Request, Msg, Resent) ->
@@ -69,9 +74,6 @@ apply2context(Key, F, A) ->
 	    {error, not_found}
     end.
 
-port_request_key(#request{key = ReqKey, socket = Socket}) ->
-    gtp_context:socket_key(Socket, ReqKey).
-
 %% TODO - MAYBE
 %%  it might be benificial to first perform the lookup and then enqueue
 %%
@@ -94,9 +96,9 @@ port_message_h(Request, #gtp{} = Msg) ->
 
 port_message_run(Request, #gtp{type = g_pdu} = Msg) ->
     port_message_p(Request, Msg);
-port_message_run(Request, Msg0) ->
+port_message_run(#request{key = ReqKey} = Request, Msg0) ->
     Msg = gtp_packet:decode_ies(Msg0),
-    case port_message(port_request_key(Request), Request, Msg, true) of
+    case port_message(ReqKey, Request, Msg, true) of
 	{error, not_found} ->
 	    port_message_p(Request, Msg);
 	Result ->
@@ -113,12 +115,6 @@ port_message_p(#request{socket = Socket} = Request, #gtp{tei = TEI} = Msg) ->
 	    Result
     end.
 
-port_message_ctx([{Handler, Server} | _], Request, Msg)
-  when is_atom(Handler), is_pid(Server) ->
-    Handler:port_message(Server, Request, Msg, false);
-port_message_ctx(_, _Request, _Msg) ->
-    throw({error, not_found}).
-
 load_class(#gtp{version = v1} = Msg) ->
     gtp_v1_c:load_class(Msg);
 load_class(#gtp{version = v2} = Msg) ->
diff --git a/src/ergw_gtp_socket.erl b/src/ergw_gtp_socket.erl
index d1f1afea..35a811ae 100644
--- a/src/ergw_gtp_socket.erl
+++ b/src/ergw_gtp_socket.erl
@@ -166,10 +166,11 @@ make_seq_id(#gtp{version = Version, seq_no = SeqNo})
 make_seq_id(_) ->
     undefined.
 
-make_request(ArrivalTS, Src, IP, Port, Msg = #gtp{version = Version, type = Type}, Socket, Info) ->
+make_request(ArrivalTS, Src, IP, Port, Msg = #gtp{version = Version, type = Type},
+	     #socket{name = SocketName} = Socket, Info) ->
     SeqId = make_seq_id(Msg),
     #request{
-       key = {Socket, IP, Port, Type, SeqId},
+       key = {request, {SocketName, IP, Port, Type, SeqId}},
        socket = Socket,
        info = Info,
        src = Src,
diff --git a/src/ergw_pfcp.erl b/src/ergw_pfcp.erl
index f797fa2e..5f82edee 100644
--- a/src/ergw_pfcp.erl
+++ b/src/ergw_pfcp.erl
@@ -158,7 +158,7 @@ outer_header_removal(v6) ->
     #outer_header_removal{header = 'GTP-U/UDP/IPv6'}.
 
 ctx_teid_key(#pfcp_ctx{name = Name}, TEI) ->
-    {Name, {teid, 'gtp-u', TEI}}.
+    #socket_teid_key{name = Name, type = 'gtp-u', teid = TEI}.
 
 up_inactivity_timer(#pfcp_ctx{up_inactivity_timer = Timer})
   when is_integer(Timer) ->
diff --git a/src/ergw_pfcp_context.erl b/src/ergw_pfcp_context.erl
index f59faecb..96e90120 100644
--- a/src/ergw_pfcp_context.erl
+++ b/src/ergw_pfcp_context.erl
@@ -755,7 +755,7 @@ make_pctx_bearer_key(_, _, _, Keys) ->
     Keys.
 
 make_pctx_keys(Bearer, #pfcp_ctx{seid = #seid{cp = SEID}} = PCtx) ->
-    maps:fold(make_pctx_bearer_key(_, _, PCtx, _), [{seid, SEID}], Bearer).
+    maps:fold(make_pctx_bearer_key(_, _, PCtx, _), [#seid_key{seid = SEID}], Bearer).
 
 register_ctx_ids(Handler, Bearer, PCtx) ->
     Keys = make_pctx_keys(Bearer, PCtx),
diff --git a/src/ergw_sx_node.erl b/src/ergw_sx_node.erl
index 6422a07b..72e221b3 100644
--- a/src/ergw_sx_node.erl
+++ b/src/ergw_sx_node.erl
@@ -183,7 +183,7 @@ init([Parent, Node, NodeSelect, IP4, IP6, NotifyUp]) ->
 
     RegKeys =
 	[gtp_context:socket_teid_key(Socket, TEI),
-	 {seid, SEID}],
+	 #seid_key{seid = SEID}],
     gtp_context_reg:register(RegKeys, ?MODULE, self()),
 
     Nodes = setup:get_env(ergw, nodes, #{}),
diff --git a/src/gtp_context.erl b/src/gtp_context.erl
index 6d0821a6..c878ffc8 100644
--- a/src/gtp_context.erl
+++ b/src/gtp_context.erl
@@ -28,7 +28,7 @@
 	 validate_options/3,
 	 validate_option/2,
 	 generic_error/3,
-	 socket_key/2, socket_teid_key/2]).
+	 context_key/2, socket_teid_key/2]).
 -export([usage_report_to_accounting/1,
 	 collect_charging_events/2]).
 
@@ -167,7 +167,7 @@ collect_charging_events(OldS, NewS) ->
 %% preexisting context should be deleted locally. This function does that.
 terminate_colliding_context(#tunnel{socket = Socket}, #context{context_id = Id})
   when Id /= undefined ->
-    case gtp_context_reg:lookup(socket_key(Socket, Id)) of
+    case gtp_context_reg:lookup(context_key(Socket, Id)) of
 	{?MODULE, Server} when is_pid(Server) ->
 	    gtp_context:terminate_context(Server);
 	_ ->
@@ -276,15 +276,15 @@ sx_report(Server, Report) ->
 port_message(Request, #gtp{version = v2, type = MsgType, tei = 0} = Msg)
   when MsgType == change_notification_request;
        MsgType == change_notification_response ->
-    Keys = gtp_v2_c:get_msg_keys(Msg),
-    ergw_context:port_message(Keys, Request, Msg);
+    Id = gtp_v2_c:get_context_id(Msg),
+    ergw_context:port_message(Id, Request, Msg);
 
 %% same as above for GTPv2
 port_message(Request, #gtp{version = v1, type = MsgType, tei = 0} = Msg)
   when MsgType == ms_info_change_notification_request;
        MsgType == ms_info_change_notification_response ->
-    Keys = gtp_v1_c:get_msg_keys(Msg),
-    ergw_context:port_message(Keys, Request, Msg);
+    Id = gtp_v1_c:get_context_id(Msg),
+    ergw_context:port_message(Id, Request, Msg);
 
 port_message(#request{socket = Socket, info = Info} = Request,
 		 #gtp{version = Version, tei = 0} = Msg) ->
@@ -733,11 +733,11 @@ generic_error(#request{socket = Socket} = Request,
 %%% Internal functions
 %%%===================================================================
 
-register_request(Handler, Server, #request{key = ReqKey, socket = Socket}) ->
-    gtp_context_reg:register([socket_key(Socket, ReqKey)], Handler, Server).
+register_request(Handler, Server, #request{key = ReqKey}) ->
+    gtp_context_reg:register([ReqKey], Handler, Server).
 
-unregister_request(#request{key = ReqKey, socket = Socket}) ->
-    gtp_context_reg:unregister([socket_key(Socket, ReqKey)], ?MODULE, self()).
+unregister_request(#request{key = ReqKey}) ->
+    gtp_context_reg:unregister([ReqKey], ?MODULE, self()).
 
 get_handler_if(Socket, #gtp{version = v1} = Msg) ->
     gtp_v1_c:get_handler(Socket, Msg);
@@ -794,7 +794,7 @@ context2keys(#tunnel{socket = Socket} = LeftTunnel, Bearer,
 	     #context{apn = APN, context_id = ContextId}) ->
     ordsets:from_list(
       tunnel2keys(LeftTunnel)
-      ++ [socket_key(Socket, ContextId) || ContextId /= undefined]
+      ++ [context_key(Socket, ContextId) || ContextId /= undefined]
       ++ maps:fold(bsf_keys(APN, _, _, _), [], Bearer)).
 
 tunnel2keys(Tunnel) ->
@@ -808,10 +808,8 @@ bsf_keys(APN, _, #bearer{vrf = VRF, local = #ue_ip{v4 = IPv4, v6 = IPv6}}, Keys)
 bsf_keys(_, _, _, Keys) ->
     Keys.
 
-socket_key(#socket{name = Name}, Key) ->
-    {Name, Key};
-socket_key({Name, _}, Key) ->
-    {Name, Key}.
+context_key(#socket{name = Name}, Id) ->
+    #context_key{socket = Name, id = Id}.
 
 tunnel_key(local, #tunnel{socket = Socket, local = #fq_teid{teid = TEID}}) ->
     socket_teid_key(Socket, TEID);
@@ -823,7 +821,7 @@ socket_teid_key(#socket{type = Type} = Socket, TEI) ->
     socket_teid_key(Socket, Type, TEI).
 
 socket_teid_key(#socket{name = Name}, Type, TEI) ->
-    {Name, {teid, Type, TEI}}.
+    #socket_teid_key{name = Name, type = Type, teid = TEI}.
 
 %%====================================================================
 %% Experimental Trigger Support
diff --git a/src/gtp_context_reg.erl b/src/gtp_context_reg.erl
index 91920fd6..f1868f1b 100644
--- a/src/gtp_context_reg.erl
+++ b/src/gtp_context_reg.erl
@@ -15,7 +15,6 @@
 -export([start_link/0]).
 -export([register/3, register_new/3, update/4, unregister/3,
 	 lookup/1, select/1,
-	 match_keys/2,
 	 await_unreg/1]).
 -export([all/0]).
 
@@ -50,19 +49,6 @@ lookup(Key) when is_tuple(Key) ->
 select(Key) ->
     ets:select(?SERVER, [{{Key, '$1'},[],['$1']}]).
 
-match_key(#socket{name = Name}, Key) ->
-    select({Name, Key}).
-
-match_keys(_, []) ->
-    throw({error, not_found});
-match_keys(Socket, [H|T]) ->
-    case match_key(Socket, H) of
-	[_|_] = Match ->
-	    Match;
-	_ ->
-	    match_keys(Socket, T)
-    end.
-
 register(Keys, Handler, Pid)
   when is_list(Keys), is_atom(Handler), is_pid(Pid) ->
     gen_server:call(?SERVER, {register, Keys, Handler, Pid}).
diff --git a/src/gtp_v1_c.erl b/src/gtp_v1_c.erl
index 1b3d08bd..242cd12f 100644
--- a/src/gtp_v1_c.erl
+++ b/src/gtp_v1_c.erl
@@ -16,7 +16,7 @@
 	 build_echo_request/0,
 	 validate_teid/2,
 	 type/0, port/0,
-	 get_msg_keys/1, update_context_id/2,
+	 get_context_id/1, update_context_id/2,
 	 get_cause/1, get_common_flags/1,
 	 find_sender_teid/1,
 	 load_class/1]).
@@ -233,7 +233,7 @@ get_common_flags(IEs) ->
 get_ext_common_flags(IEs) ->
     get_element(?'Extended Common Flags', IEs, #extended_common_flags.flags, []).
 
-get_context_id(IEs) ->
+get_context_id(#gtp{version = v1, ie = IEs}) ->
     NSAPI = get_element(?'NSAPI', IEs, #nsapi.nsapi, '_'),
     UIMSI = proplists:get_bool('Unauthenticated IMSI', get_ext_common_flags(IEs)),
     %% order of key selection, first match terminates:
@@ -251,16 +251,8 @@ get_context_id(IEs) ->
 	    undefined
     end.
 
-get_msg_keys(#gtp{version = v1, ie = IEs}) ->
-    case get_context_id(IEs) of
-	undefined ->
-	    [];
-	Id ->
-	    [Id]
-    end.
-
-update_context_id(#gtp{version = v1, ie = IEs}, Context) ->
-    case get_context_id(IEs) of
+update_context_id(Msg, Context) ->
+    case get_context_id(Msg) of
 	{_, _, NSAPI} = Id when is_integer(NSAPI) ->
 	    Context#context{context_id = Id};
 	_ ->
diff --git a/src/gtp_v2_c.erl b/src/gtp_v2_c.erl
index 2449dc28..c5a126ca 100644
--- a/src/gtp_v2_c.erl
+++ b/src/gtp_v2_c.erl
@@ -16,7 +16,7 @@
 	 build_echo_request/0,
 	 validate_teid/2,
 	 type/0, port/0,
-	 get_msg_keys/1, update_context_id/2,
+	 get_context_id/1, update_context_id/2,
 	 get_cause/1, get_indication_flags/1,
 	 find_sender_teid/1,
 	 load_class/1]).
@@ -258,7 +258,7 @@ get_context_ebi(#{?'Bearer Contexts' :=
 get_context_ebi(_) ->
     '_'.
 
-get_context_id(IEs) ->
+get_context_id(#gtp{version = v2, ie = IEs}) ->
     EBI = get_context_ebi(IEs),
     UIMSI = proplists:get_bool('UIMSI', get_indication_flags(IEs)),
     case {UIMSI, IEs} of
@@ -270,16 +270,8 @@ get_context_id(IEs) ->
 	    undefined
     end.
 
-get_msg_keys(#gtp{version = v2, ie = IEs}) ->
-    case get_context_id(IEs) of
-	undefined ->
-	    [];
-	Id ->
-	    [Id]
-    end.
-
-update_context_id(#gtp{version = v2, ie = IEs}, Context) ->
-    case get_context_id(IEs) of
+update_context_id(Msg, Context) ->
+    case get_context_id(Msg) of
 	{_, _, EBI} = Id when is_integer(EBI) ->
 	    Context#context{context_id = Id};
 	_Other ->
diff --git a/test/ergw_test_lib.erl b/test/ergw_test_lib.erl
index 7bc30321..d1de8eab 100644
--- a/test/ergw_test_lib.erl
+++ b/test/ergw_test_lib.erl
@@ -43,6 +43,7 @@
 -export([init_ets/1]).
 
 -include("ergw_test_lib.hrl").
+-include_lib("stdlib/include/ms_transform.hrl").
 -include_lib("kernel/include/logger.hrl").
 -include_lib("common_test/include/ct.hrl").
 -include_lib("gtplib/include/gtp_packet.hrl").
@@ -544,7 +545,8 @@ add_cfg_value([H | T], Value, Config) ->
 %%%===================================================================
 
 outstanding_requests() ->
-    ets:match_object(gtp_context_reg, {{'_', {'_', '_', '_', '_', '_'}}, '_'}).
+    Ms = ets:fun2ms(fun({Key, _} = Obj) when element(1, Key) == 'request' -> Obj end),
+    ets:select(gtp_context_reg, Ms).
 
 wait4tunnels(Cnt) ->
     case [X || X = #{tunnels := T} <- ergw_api:peer(all), T /= 0] of
diff --git a/test/ergw_test_sx_up.erl b/test/ergw_test_sx_up.erl
index 174d191e..28523f9d 100644
--- a/test/ergw_test_sx_up.erl
+++ b/test/ergw_test_sx_up.erl
@@ -172,7 +172,8 @@ handle_call({send, Msg}, _From,
 	    #state{gtp = GtpSocket, cp_ip = IP, up_ip = UpIP} = State)
   when is_binary(Msg) ->
     {ok, SxPid} = ergw_sx_node_reg:lookup(ergw_inet:bin2ip(UpIP)),
-    [[SxTEI]] = ets:match(gtp_context_reg, {{'cp-socket',{teid,'gtp-u','$1'}},{'_',SxPid}}),
+    TEIDMatch = #socket_teid_key{name = 'cp-socket', type = 'gtp-u', teid = '$1', _ = '_'},
+    [[SxTEI]] = ets:match(gtp_context_reg, {TEIDMatch, {'_',SxPid}}),
     BinMsg = gtp_packet:encode(#gtp{version = v1, type = g_pdu, tei = SxTEI, ie = Msg}),
     ok = gen_udp:send(GtpSocket, IP, ?GTP1u_PORT, BinMsg),
     {reply, ok, State};
diff --git a/test/ggsn_SUITE.erl b/test/ggsn_SUITE.erl
index 4ea00511..4e6d9759 100644
--- a/test/ggsn_SUITE.erl
+++ b/test/ggsn_SUITE.erl
@@ -973,9 +973,10 @@ path_failure() ->
     [{doc, "Check that Create PDP Context works and "
       "that a path failure (Echo timeout) terminates the session"}].
 path_failure(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, CtxPid} = gtp_context_reg:lookup(CtxKey),
     #{left_tunnel := #tunnel{socket = CSocket}} = gtp_context:info(CtxPid),
 
     ClientIP = proplists:get_value(client_ip, Config),
@@ -1414,10 +1415,11 @@ delete_pdp_context_requested() ->
     [{doc, "Check GGSN initiated Delete PDP Context"}].
 delete_pdp_context_requested(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
 
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -1449,10 +1451,11 @@ delete_pdp_context_requested_resend() ->
     [{doc, "Check resend of GGSN initiated Delete PDP Context"}].
 delete_pdp_context_requested_resend(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
 
     {_, _, _} = create_pdp_context(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -1717,6 +1720,7 @@ gy_validity_timer(Config) ->
 simple_aaa() ->
     [{doc, "Check simple session with RADIOS/DIAMETER over (S)Gi"}].
 simple_aaa(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
     Interim = rand:uniform(1800) + 1800,
     AAAReply = #{'Acct-Interim-Interval' => Interim},
 
@@ -1733,7 +1737,7 @@ simple_aaa(Config) ->
 
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
 
@@ -1826,6 +1830,7 @@ simple_aaa(Config) ->
 simple_ofcs() ->
     [{doc, "Check simple session with DIAMETER Rf"}].
 simple_ofcs(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
     Interim = rand:uniform(1800) + 1800,
     AAAReply = #{'Acct-Interim-Interval' => [Interim]},
 
@@ -1842,7 +1847,7 @@ simple_ofcs(Config) ->
 
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
 
@@ -1975,9 +1980,10 @@ simple_ofcs(Config) ->
 simple_ocs() ->
     [{doc, "Test Gy a simple interaction"}].
 simple_ocs(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
 
@@ -2164,10 +2170,11 @@ gy_ccr_asr_overlap() ->
     [{doc, "Test that ASR is answered when it arrives during CCR-T"}].
 gy_ccr_asr_overlap(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
 
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     #{'Session' := Session} = gtp_context:info(Server),
@@ -2296,9 +2303,10 @@ volume_threshold(Config) ->
 gx_rar_gy_interaction() ->
     [{doc, "Check that a Gx RAR triggers a Gy request"}].
 gx_rar_gy_interaction(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     {ok, Session} = gtp_context:test_cmd(Server, session),
@@ -2356,10 +2364,11 @@ gx_asr() ->
     [{doc, "Check that ASR on Gx terminates the session"}].
 gx_asr(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
 
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     ResponseFun = fun(_, _, _, _) -> ok end,
@@ -2382,9 +2391,10 @@ gx_asr(Config) ->
 gx_rar() ->
     [{doc, "Check that RAR on Gx changes the session"}].
 gx_rar(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     #{'Session' := Session} = gtp_context:info(Server),
@@ -2491,10 +2501,11 @@ gy_asr() ->
     [{doc, "Check that ASR on Gy terminates the session"}].
 gy_asr(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
 
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     ResponseFun = fun(_, _, _, _) -> ok end,
@@ -2627,6 +2638,7 @@ gtp_idle_timeout(Config) ->
 up_inactivity_timer() ->
     [{doc, "Test expiry of the User Plane Inactivity Timer"}].
 up_inactivity_timer(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
     Interim = rand:uniform(1800) + 1800,
     AAAReply = #{'Acct-Interim-Interval' => Interim},
 
@@ -2644,7 +2656,8 @@ up_inactivity_timer(Config) ->
 	   end),
 
     create_pdp_context(Config),
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
     [SER|_] = lists:filter(
diff --git a/test/ggsn_proxy_SUITE.erl b/test/ggsn_proxy_SUITE.erl
index 51a7d1ce..b61de30d 100644
--- a/test/ggsn_proxy_SUITE.erl
+++ b/test/ggsn_proxy_SUITE.erl
@@ -1013,9 +1013,10 @@ ggsn_broken_recovery() ->
     [{doc, "Check that Create PDP Context Request works and "
            "that a GGSN Restart terminates the session"}].
 ggsn_broken_recovery(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, CtxPid} = gtp_context_reg:lookup(CtxKey),
     #{right_tunnel := #tunnel{socket = CSocket}} = gtp_context:info(CtxPid),
 
     FinalGSN = proplists:get_value(final_gsn, Config),
@@ -1049,10 +1050,12 @@ path_failure_to_ggsn() ->
       "that a path failure (Echo timeout) terminates the session"}].
 path_failure_to_ggsn(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
 
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, CtxPid} = gtp_context_reg:lookup(CtxKey),
     #{right_tunnel := #tunnel{socket = CSocket}} = gtp_context:info(CtxPid),
 
     FinalGSN = proplists:get_value(final_gsn, Config),
@@ -1077,7 +1080,7 @@ path_failure_to_ggsn(Config) ->
     ct:sleep(100),
     delete_pdp_context(not_found, GtpC),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
     %% killing the GGSN context
     exit(Server, kill),
@@ -1103,12 +1106,13 @@ path_failure_to_ggsn_and_restore() ->
       "and is later restored with a valid echo"}].
 path_failure_to_ggsn_and_restore(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
 
     lists:foreach(fun({_, Pid, _}) -> gtp_path:stop(Pid) end, gtp_path_reg:all()),
 
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, CtxPid} = gtp_context_reg:lookup(CtxKey),
     #{right_tunnel := #tunnel{socket = CSocket}} = gtp_context:info(CtxPid),
 
     FinalGSN = proplists:get_value(final_gsn, Config),
@@ -1167,9 +1171,10 @@ path_failure_to_sgsn() ->
     [{doc, "Check that Create PDP Context works and "
       "that a path failure (Echo timeout) terminates the session"}].
 path_failure_to_sgsn(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, CtxPid} = gtp_context_reg:lookup(CtxKey),
     #{left_tunnel := #tunnel{socket = CSocket}} = gtp_context:info(CtxPid),
 
     ClientIP = proplists:get_value(client_ip, Config),
@@ -1280,7 +1285,8 @@ create_lb_multi_context(Config) ->
 	    {GtpC1, _, _} = create_pdp_context(#{apn => ?'APN-LB-1'}, GtpC0),
 
 	    CntS = lists:foldl(
-		     fun({{_,{teid,'gtp-c',{fq_teid, PeerIP,_}}},_}, M)
+		     fun({#socket_teid_key{
+			     type = 'gtp-c', teid = #fq_teid{ip = PeerIP}}, _}, M)
 			   when is_map_key(PeerIP, M) ->
 			     maps:update_with(PeerIP, fun(C) -> C + 1 end, 1, M);
 			(_, M) -> M
@@ -1338,7 +1344,9 @@ one_lb_node_down(Config) ->
     GtpCs0 = make_gtp_contexts(?NUM_OF_CLIENTS, Config),
     GtpCs1 = lists:map(fun(GtpC0) -> create_pdp_context(random, GtpC0) end, GtpCs0),
 
-    PgwFqTeids = [X || {{_,{teid,'gtp-c',{fq_teid, PeerIP,_}}},_} =
+
+
+    PgwFqTeids = [X || {#socket_teid_key{type = 'gtp-c', teid = #fq_teid{ip = PeerIP}}, _} =
 			   X <- gtp_context_reg:all(), PeerIP == DownGSN],
     ?match(0, length(PgwFqTeids)), % Check no connection to down peer
 
@@ -1390,12 +1398,13 @@ create_pdp_context_request_resend(Config) ->
 create_pdp_context_proxy_request_resend() ->
     [{doc, "Check that the proxy does not send the Create PDP Context Request multiple times"}].
 create_pdp_context_proxy_request_resend(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
     GtpC = gtp_context(Config),
     Request = make_request(create_pdp_context_request, simple, GtpC),
 
     ?equal({error,timeout}, send_recv_pdu(GtpC, Request, 2 * 1000, error)),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     %% killing the proxy PGW context
@@ -1411,13 +1420,14 @@ create_pdp_context_request_timeout() ->
     [{doc, "Check that the proxy does shutdown the context on timeout"}].
 create_pdp_context_request_timeout(Config) ->
     %% logger:set_primary_config(level, debug),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
 
     GtpC = gtp_context(Config),
     Request = make_request(create_pdp_context_request, simple, GtpC),
 
     ?equal({error,timeout}, send_recv_pdu(GtpC, Request, 2 * 1000, error)),
 
-    ?equal(undefined, gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}})),
+    ?equal(undefined, gtp_context_reg:lookup(CtxKey)),
     ?match(1, meck:num_calls(ggsn_gn, handle_request, '_')),
 
     wait4tunnels(?TIMEOUT),
@@ -1445,8 +1455,11 @@ delete_pdp_context_request_timeout() ->
     [{doc, "Check that a Delete PDP Context Request terminates the "
            "proxy session even when the final GSN fails"}].
 delete_pdp_context_request_timeout(Config) ->
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
+
     {GtpC, _, _} = create_pdp_context(Config),
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
 
     Request = make_request(delete_pdp_context_request, simple, GtpC),
@@ -1486,10 +1499,12 @@ error_indication_ggsn2sgsn() ->
     [{doc, "Check the a GTP-U error indication terminates the session"}].
 error_indication_ggsn2sgsn(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
 
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, CtxPid} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(CtxPid),
     #{bearer := #{right := RightBearer}} = gtp_context:info(CtxPid),
 
@@ -1503,7 +1518,7 @@ error_indication_ggsn2sgsn(Config) ->
     ct:sleep(100),
     delete_pdp_context(not_found, GtpC),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
     %% killing the GGSN context
     exit(Server, kill),
@@ -1552,8 +1567,10 @@ request_fast_resend(Config) ->
 update_pdp_context_request_ra_update() ->
     [{doc, "Check Update PDP Context with Routing Area Update"}].
 update_pdp_context_request_ra_update(Config) ->
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
+
     {GtpC1, _, _} = create_pdp_context(Config),
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, CtxPid} = gtp_context_reg:lookup(RemoteCtxKey),
     #{left_tunnel := LeftTunnel1, bearer := #{left := LeftBearer1}} = gtp_context:info(CtxPid),
 
     {GtpC2, _, _} = update_pdp_context(ra_update, GtpC1),
@@ -1578,11 +1595,14 @@ update_pdp_context_request_ra_update(Config) ->
 update_pdp_context_request_tei_update() ->
     [{doc, "Check Update PDP Context with TEID update (e.g. SGSN change)"}].
 update_pdp_context_request_tei_update(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
+
     {GtpC1, _, _} = create_pdp_context(Config),
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, CtxPid} = gtp_context_reg:lookup(RemoteCtxKey),
     #{left_tunnel := LeftTunnel1, bearer := #{left := LeftBearer1}} = gtp_context:info(CtxPid),
 
-    {_Handler, ProxyCtxPid} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, ProxyCtxPid} = gtp_context_reg:lookup(CtxKey),
     #{right_tunnel := RightTunnel1} = gtp_context:info(ProxyCtxPid),
     ProxyRegKey1 = gtp_context:tunnel_key(local, RightTunnel1),
     ?match({gtp_context, ProxyCtxPid}, gtp_context_reg:lookup(ProxyRegKey1)),
@@ -1627,6 +1647,8 @@ update_pdp_context_request_tei_update(Config) ->
 update_pdp_context_request_broken_recovery() ->
     [{doc, "Check Update PDP Context where the response includes an invalid recovery"}].
 update_pdp_context_request_broken_recovery(Config) ->
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
+
     ok = meck:expect(gtp_context, send_response,
 		     fun(ReqKey, Request, {update_pdp_context_response, TEID, IEs0})->
 			     IEs = [#recovery{restart_counter = 0},
@@ -1645,7 +1667,7 @@ update_pdp_context_request_broken_recovery(Config) ->
 			     meck:passthrough([ReqKey, Request, Response])
 		     end),
     {GtpC1, _, _} = create_pdp_context(Config),
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, CtxPid} = gtp_context_reg:lookup(RemoteCtxKey),
     #{left_tunnel := LeftTunnel1, bearer := #{left := LeftBearer1}} = gtp_context:info(CtxPid),
 
     {GtpC2, _, _} = update_pdp_context(simple, GtpC1),
@@ -1805,10 +1827,11 @@ delete_pdp_context_requested() ->
     [{doc, "Check GGSN initiated Delete PDP Context"}].
 delete_pdp_context_requested(Config) ->
     Cntl = whereis(gtpc_client_server),
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
 
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -1839,10 +1862,11 @@ delete_pdp_context_requested_resend() ->
     [{doc, "Check resend of GGSN initiated Delete PDP Context"}].
 delete_pdp_context_requested_resend(Config) ->
     Cntl = whereis(gtpc_client_server),
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
 
     {_, _, _} = create_pdp_context(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -1871,10 +1895,11 @@ delete_pdp_context_requested_invalid_teid() ->
     [{doc, "Check error response of GGSN initiated Delete PDP Context with invalid TEID"}].
 delete_pdp_context_requested_invalid_teid(Config) ->
     Cntl = whereis(gtpc_client_server),
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
 
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -1908,8 +1933,9 @@ delete_pdp_context_requested_late_response(Config) ->
     Cntl = whereis(gtpc_client_server),
 
     {GtpC, _, _} = create_pdp_context(Config),
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
 
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -1944,10 +1970,11 @@ ggsn_update_pdp_context_request() ->
      {timetrap,{seconds,60}}].
 ggsn_update_pdp_context_request(Config) ->
     Cntl = whereis(gtpc_client_server),
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
 
     {GtpC, _, _} = create_pdp_context(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -1978,10 +2005,12 @@ ggsn_update_pdp_context_request(Config) ->
 create_pdp_context_overload() ->
     [{doc, "Check that the overload protection works"}].
 create_pdp_context_overload(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
+
     create_pdp_context(overload, Config),
 
     ct:sleep(10),
-    ?equal(undefined, gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}})),
+    ?equal(undefined, gtp_context_reg:lookup(CtxKey)),
 
     ?equal([], outstanding_requests()),
     meck_validate(Config),
diff --git a/test/pgw_SUITE.erl b/test/pgw_SUITE.erl
index 5a0391c8..cde0f703 100644
--- a/test/pgw_SUITE.erl
+++ b/test/pgw_SUITE.erl
@@ -1247,9 +1247,11 @@ path_failure() ->
     [{doc, "Check that Create Session Request works and "
       "that a path failure (Echo timeout) terminates the session"}].
 path_failure(Config) ->
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
+
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, CtxPid} = gtp_context_reg:lookup(CtxKey),
     #{left_tunnel := #tunnel{socket = CSocket}} = gtp_context:info(CtxPid),
 
     ClientIP = proplists:get_value(client_ip, Config),
@@ -2162,10 +2164,11 @@ delete_bearer_request() ->
      {timetrap,{seconds,60}}].
 delete_bearer_request(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -2240,10 +2243,11 @@ delete_bearer_request_resend() ->
      {timetrap,{seconds,60}}].
 delete_bearer_request_resend(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
 
     {_, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -2824,6 +2828,7 @@ gy_validity_timer(Config) ->
 simple_aaa() ->
     [{doc, "Check simple session with RADIOS/DIAMETER over (S)Gi"}].
 simple_aaa(Config) ->
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
     Interim = rand:uniform(1800) + 1800,
     AAAReply = #{'Acct-Interim-Interval' => Interim},
 
@@ -2840,7 +2845,7 @@ simple_aaa(Config) ->
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
 
@@ -2935,6 +2940,7 @@ simple_aaa(Config) ->
 simple_ofcs() ->
     [{doc, "Check simple session with DIAMETER Rf"}].
 simple_ofcs(Config) ->
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
     Interim = rand:uniform(1800) + 1800,
     AAAReply = #{'Acct-Interim-Interval' => [Interim]},
 
@@ -2951,7 +2957,7 @@ simple_ofcs(Config) ->
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
 
@@ -3086,6 +3092,7 @@ simple_ofcs(Config) ->
 ofcs_no_interim() ->
     [{doc, "Check that OFCS reporting also works without Acct-Interim-Interval"}].
 ofcs_no_interim(Config) ->
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
     AAAReply = #{},
 
     ok = meck:expect(ergw_aaa_session, invoke,
@@ -3101,7 +3108,7 @@ ofcs_no_interim(Config) ->
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, _} = gtp_context:test_cmd(Server, pfcp_ctx),
 
@@ -3252,9 +3259,11 @@ secondary_rat_usage_data_report(Config) ->
 simple_ocs() ->
     [{doc, "Test Gy a simple interaction"}].
 simple_ocs(Config) ->
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
+
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
 
@@ -3439,6 +3448,7 @@ split_charging1() ->
     [{doc, "Used different Rating-Groups for Online and Offline charging, "
       "without catch all PCC rules/RG"}].
 split_charging1(Config) ->
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
     Interim = rand:uniform(1800) + 1800,
     AAAReply = #{'Acct-Interim-Interval' => [Interim]},
 
@@ -3455,7 +3465,7 @@ split_charging1(Config) ->
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
 
@@ -3755,6 +3765,7 @@ split_charging2() ->
     [{doc, "Used different Rating-Groups for Online and Offline charging, "
       "with catch all PCC rules/RG"}].
 split_charging2(Config) ->
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
     Interim = rand:uniform(1800) + 1800,
     AAAReply = #{'Acct-Interim-Interval' => [Interim]},
 
@@ -3771,7 +3782,7 @@ split_charging2(Config) ->
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
 
@@ -4157,6 +4168,7 @@ aa_pool_select_fail(Config) ->
 tariff_time_change() ->
     [{doc, "Check Rf and Gy action on Tariff-Time-Change"}].
 tariff_time_change(Config) ->
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
     Interim = rand:uniform(1800) + 1800,
     AAAReply = #{'Acct-Interim-Interval' => [Interim]},
 
@@ -4179,7 +4191,7 @@ tariff_time_change(Config) ->
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
 
@@ -4401,10 +4413,11 @@ gy_ccr_asr_overlap() ->
     [{doc, "Test that ASR is answered when it arrives during CCR-T"}].
 gy_ccr_asr_overlap(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     #{'Session' := Session} = gtp_context:info(Server),
@@ -4535,9 +4548,11 @@ volume_threshold(Config) ->
 gx_rar_gy_interaction() ->
     [{doc, "Check that a Gx RAR triggers a Gy request"}].
 gx_rar_gy_interaction(Config) ->
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
+
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     {ok, Session} = gtp_context:test_cmd(Server, session),
@@ -4730,10 +4745,11 @@ gx_asr() ->
     [{doc, "Check that ASR on Gx terminates the session"}].
 gx_asr(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     ResponseFun = fun(_, _, _, _) -> ok end,
@@ -4756,9 +4772,11 @@ gx_asr(Config) ->
 gx_rar() ->
     [{doc, "Check that RAR on Gx changes the session"}].
 gx_rar(Config) ->
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
+
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     #{'Session' := Session} = gtp_context:info(Server),
@@ -4883,10 +4901,11 @@ gy_asr() ->
     [{doc, "Check that ASR on Gy terminates the session"}].
 gy_asr(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     ResponseFun = fun(_, _, _, _) -> ok end,
@@ -5174,6 +5193,7 @@ gtp_idle_timeout(Config) ->
 up_inactivity_timer() ->
     [{doc, "Test expiry of the User Plane Inactivity Timer"}].
 up_inactivity_timer(Config) ->
+    CtxKey = #context_key{socket = 'irx-socket', id = {imsi, ?'IMSI', 5}},
     Interim = rand:uniform(1800) + 1800,
     AAAReply = #{'Acct-Interim-Interval' => Interim},
 
@@ -5191,7 +5211,8 @@ up_inactivity_timer(Config) ->
 	   end),
 
     create_session(Config),
-    {_Handler, Server} = gtp_context_reg:lookup({'irx-socket', {imsi, ?'IMSI', 5}}),
+
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
     [SER|_] = lists:filter(
diff --git a/test/pgw_proxy_SUITE.erl b/test/pgw_proxy_SUITE.erl
index 11a4377b..9ee89fec 100644
--- a/test/pgw_proxy_SUITE.erl
+++ b/test/pgw_proxy_SUITE.erl
@@ -1113,10 +1113,12 @@ path_failure_to_pgw() ->
       "that a path failure (Echo timeout) terminates the session"}].
 path_failure_to_pgw(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, CtxPid} = gtp_context_reg:lookup(CtxKey),
     #{right_tunnel := #tunnel{socket = CSocket}} = gtp_context:info(CtxPid),
 
     FinalGSN = proplists:get_value(final_gsn, Config),
@@ -1141,7 +1143,7 @@ path_failure_to_pgw(Config) ->
     ct:sleep(100),
     delete_session(not_found, GtpC),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
     %% killing the PGW context
     exit(Server, kill),
@@ -1165,12 +1167,13 @@ path_failure_to_pgw_and_restore() ->
       "and is later restored with a valid echo"}].
 path_failure_to_pgw_and_restore(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
 
     lists:foreach(fun({_, Pid, _}) -> gtp_path:stop(Pid) end, gtp_path_reg:all()),
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, CtxPid} = gtp_context_reg:lookup(CtxKey),
     #{right_tunnel := #tunnel{socket = CSocket}} = gtp_context:info(CtxPid),
 
     FinalGSN = proplists:get_value(final_gsn, Config),
@@ -1230,9 +1233,11 @@ path_failure_to_sgw() ->
     [{doc, "Check that Create Session Request works and "
       "that a path failure (Echo timeout) terminates the session"}].
 path_failure_to_sgw(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
+
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, CtxPid} = gtp_context_reg:lookup(CtxKey),
     #{left_tunnel := #tunnel{socket = CSocket}} = gtp_context:info(CtxPid),
 
     ClientIP = proplists:get_value(client_ip, Config),
@@ -1609,7 +1614,8 @@ create_lb_multi_session(Config) ->
 	    {GtpC1, _, _} = create_session(#{apn => ?'APN-LB-1'}, GtpC0),
 
 	    CntS = lists:foldl(
-		     fun({{_,{teid,'gtp-c',{fq_teid, PeerIP,_}}},_}, M)
+		     fun({#socket_teid_key{
+			     type = 'gtp-c', teid = #fq_teid{ip = PeerIP}}, _}, M)
 			   when is_map_key(PeerIP, M) ->
 			     maps:update_with(PeerIP, fun(C) -> C + 1 end, 1, M);
 			(_, M) -> M
@@ -1667,7 +1673,7 @@ one_lb_node_down(Config) ->
     GtpCs0 = make_gtp_contexts(?NUM_OF_CLIENTS, Config),
     GtpCs1 = lists:map(fun(GtpC0) -> create_session(random, GtpC0) end, GtpCs0),
 
-    PgwFqTeids = [X || {{_,{teid,'gtp-c',{fq_teid, PeerIP,_}}},_} =
+    PgwFqTeids = [X || {#socket_teid_key{type = 'gtp-c', teid = #fq_teid{ip = PeerIP}}, _} =
 			   X <- gtp_context_reg:all(), PeerIP == DownGSN],
     ?match(0, length(PgwFqTeids)), % Check no connection to down peer
 
@@ -1718,12 +1724,14 @@ create_session_request_resend(Config) ->
 create_session_proxy_request_resend() ->
     [{doc, "Check that the proxy does not send the Create Session Request multiple times"}].
 create_session_proxy_request_resend(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
+
     GtpC = gtp_context(Config),
     Request = make_request(create_session_request, simple, GtpC),
 
     ?equal({error,timeout}, send_recv_pdu(GtpC, Request, 2 * 1000, error)),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     %% killing the proxy PGW context
@@ -1770,8 +1778,11 @@ delete_session_request_timeout() ->
     [{doc, "Check that a Delete Session Request terminates the "
            "proxy session even when the final GSN fails"}].
 delete_session_request_timeout(Config) ->
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
+
     {GtpC, _, _} = create_session(Config),
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
 
     Request = make_request(delete_session_request, simple, GtpC),
@@ -1811,10 +1822,12 @@ error_indication_pgw2sgw() ->
     [{doc, "Check the a GTP-U error indication terminates the session"}].
 error_indication_pgw2sgw(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, CtxPid} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(CtxPid),
     #{bearer := #{right := RightBearer}} = gtp_context:info(CtxPid),
 
@@ -1828,7 +1841,7 @@ error_indication_pgw2sgw(Config) ->
     ct:sleep(100),
     delete_session(not_found, GtpC),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
     %% killing the PGW context
     exit(Server, kill),
@@ -1877,8 +1890,11 @@ request_fast_resend(Config) ->
 modify_bearer_request_ra_update() ->
     [{doc, "Check Modify Bearer Routing Area Update"}].
 modify_bearer_request_ra_update(Config) ->
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
+
     {GtpC1, _, _} = create_session(Config),
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+
+    {_Handler, CtxPid} = gtp_context_reg:lookup(RemoteCtxKey),
     #{left_tunnel := LeftTunnel1, bearer := #{left := LeftBearer1}} = gtp_context:info(CtxPid),
 
     {GtpC2, _, _} = modify_bearer(ra_update, GtpC1),
@@ -1903,11 +1919,15 @@ modify_bearer_request_ra_update(Config) ->
 modify_bearer_request_tei_update() ->
     [{doc, "Check Modify Bearer with TEID update (e.g. SGW change)"}].
 modify_bearer_request_tei_update(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
+
     {GtpC1, _, _} = create_session(Config),
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+
+    {_Handler, CtxPid} = gtp_context_reg:lookup(RemoteCtxKey),
     #{left_tunnel := LeftTunnel1, bearer := #{left := LeftBearer1}} = gtp_context:info(CtxPid),
 
-    {_Handler, ProxyCtxPid} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, ProxyCtxPid} = gtp_context_reg:lookup(CtxKey),
     #{right_tunnel := RightTunnel1} = gtp_context:info(ProxyCtxPid),
     ProxyRegKey1 = gtp_context:tunnel_key(local, RightTunnel1),
     ?match({gtp_context, ProxyCtxPid}, gtp_context_reg:lookup(ProxyRegKey1)),
@@ -2232,10 +2252,11 @@ delete_bearer_request() ->
      {timetrap,{seconds,60}}].
 delete_bearer_request(Config) ->
     Cntl = whereis(gtpc_client_server),
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -2267,10 +2288,11 @@ delete_bearer_request_resend() ->
      {timetrap,{seconds,60}}].
 delete_bearer_request_resend(Config) ->
     Cntl = whereis(gtpc_client_server),
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
 
     {_, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -2300,10 +2322,11 @@ delete_bearer_request_invalid_teid() ->
      {timetrap,{seconds,60}}].
 delete_bearer_request_invalid_teid(Config) ->
     Cntl = whereis(gtpc_client_server),
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -2336,10 +2359,11 @@ delete_bearer_request_late_response() ->
      {timetrap,{seconds,60}}].
 delete_bearer_request_late_response(Config) ->
     Cntl = whereis(gtpc_client_server),
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -2388,8 +2412,11 @@ unsupported_request(Config) ->
 interop_sgsn_to_sgw() ->
     [{doc, "Check 3GPP T 23.401, Annex D, SGSN to SGW handover"}].
 interop_sgsn_to_sgw(Config) ->
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
+
     {GtpC1, _, _} = ergw_ggsn_test_lib:create_pdp_context(Config),
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+
+    {_Handler, CtxPid} = gtp_context_reg:lookup(RemoteCtxKey),
     #{left_tunnel := LeftTunnel1, bearer := #{left := LeftBearer1}} = gtp_context:info(CtxPid),
 
     check_contexts_metric(v1, 3, 1),
@@ -2437,8 +2464,11 @@ interop_sgsn_to_sgw(Config) ->
 interop_sgw_to_sgsn() ->
     [{doc, "Check 3GPP T 23.401, Annex D, SGW to SGSN handover"}].
 interop_sgw_to_sgsn(Config) ->
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
+
     {GtpC1, _, _} = create_session(Config),
-    {_Handler, CtxPid} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+
+    {_Handler, CtxPid} = gtp_context_reg:lookup(RemoteCtxKey),
     #{left_tunnel := LeftTunnel1, bearer := #{left := LeftBearer1}} = gtp_context:info(CtxPid),
 
     check_contexts_metric(v1, 0, 0),
@@ -2488,10 +2518,11 @@ update_bearer_request() ->
      {timetrap,{seconds,60}}].
 update_bearer_request(Config) ->
     Cntl = whereis(gtpc_client_server),
+    RemoteCtxKey = #context_key{socket = 'remote-irx', id = {imsi, ?'PROXY-IMSI', 5}},
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'remote-irx', {imsi, ?'PROXY-IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(RemoteCtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -2830,7 +2861,7 @@ check_contexts_metric(Version, Cnt, Expect) ->
     [?equal({Path, Expect}, M) || {Path, _} = M <- Metrics].
 
 pfcp_seids() ->
-    lists:flatten(ets:match(gtp_context_reg, {{seid, '$1'},{ergw_sx_node, '_'}})).
+    lists:flatten(ets:match(gtp_context_reg, {#seid_key{seid = '$1'},{ergw_sx_node, '_'}})).
 
 % Set Timers for Path management
 set_path_timers(SetTimers) ->
diff --git a/test/saegw_s11_SUITE.erl b/test/saegw_s11_SUITE.erl
index af074b36..e5f972ac 100644
--- a/test/saegw_s11_SUITE.erl
+++ b/test/saegw_s11_SUITE.erl
@@ -999,10 +999,11 @@ delete_bearer_request() ->
      {timetrap,{seconds,60}}].
 delete_bearer_request(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({irx, {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -1034,10 +1035,11 @@ delete_bearer_request_resend() ->
      {timetrap,{seconds,60}}].
 delete_bearer_request_resend(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
 
     {_, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({irx, {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     Self = self(),
@@ -1238,6 +1240,7 @@ gy_validity_timer(Config) ->
 simple_aaa() ->
     [{doc, "Check simple session with RADIOS/DIAMETER over (S)Gi"}].
 simple_aaa(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
     Interim = rand:uniform(1800) + 1800,
     AAAReply = #{'Acct-Interim-Interval' => Interim},
 
@@ -1255,7 +1258,7 @@ simple_aaa(Config) ->
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
 
@@ -1346,6 +1349,7 @@ simple_aaa(Config) ->
 simple_ofcs() ->
     [{doc, "Check simple session with DIAMETER Rf"}].
 simple_ofcs(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
     Interim = rand:uniform(1800) + 1800,
     AAAReply = #{'Acct-Interim-Interval' => [Interim]},
 
@@ -1362,7 +1366,7 @@ simple_ofcs(Config) ->
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
 
@@ -1494,9 +1498,11 @@ simple_ofcs(Config) ->
 simple_ocs() ->
     [{doc, "Test Gy a simple interaction"}].
 simple_ocs(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
+
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
 
@@ -1680,10 +1686,11 @@ gy_ccr_asr_overlap() ->
     [{doc, "Test that ASR is answered when it arrives during CCR-T"}].
 gy_ccr_asr_overlap(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     #{'Session' := Session} = gtp_context:info(Server),
@@ -1810,9 +1817,11 @@ volume_threshold(Config) ->
 gx_rar_gy_interaction() ->
     [{doc, "Check that a Gx RAR triggers a Gy request"}].
 gx_rar_gy_interaction(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
+
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     {ok, Session} = gtp_context:test_cmd(Server, session),
@@ -1868,10 +1877,11 @@ gx_asr() ->
     [{doc, "Check that ASR on Gx terminates the session"}].
 gx_asr(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     ResponseFun = fun(_, _, _, _) -> ok end,
@@ -1892,10 +1902,12 @@ gx_asr(Config) ->
 gx_rar() ->
     [{doc, "Check that RAR on Gx changes the session"}].
 gx_rar(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
+
     {GtpC1, _, _} = create_session(Config),
     {GtpC2, _, _} = modify_bearer(enb_u_tei, GtpC1),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     #{'Session' := Session} = gtp_context:info(Server),
@@ -2000,10 +2012,11 @@ gy_asr() ->
     [{doc, "Check that ASR on Gy terminates the session"}].
 gy_asr(Config) ->
     Cntl = whereis(gtpc_client_server),
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
 
     {GtpC, _, _} = create_session(Config),
 
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
 
     ResponseFun = fun(_, _, _, _) -> ok end,
@@ -2126,6 +2139,7 @@ gtp_idle_timeout(Config) ->
 up_inactivity_timer() ->
     [{doc, "Test expiry of the User Plane Inactivity Timer"}].
 up_inactivity_timer(Config) ->
+    CtxKey = #context_key{socket = 'irx', id = {imsi, ?'IMSI', 5}},
     Interim = rand:uniform(1800) + 1800,
     AAAReply = #{'Acct-Interim-Interval' => Interim},
 
@@ -2143,7 +2157,8 @@ up_inactivity_timer(Config) ->
 	   end),
 
     create_session(Config),
-    {_Handler, Server} = gtp_context_reg:lookup({'irx', {imsi, ?'IMSI', 5}}),
+
+    {_Handler, Server} = gtp_context_reg:lookup(CtxKey),
     true = is_pid(Server),
     {ok, PCtx} = gtp_context:test_cmd(Server, pfcp_ctx),
     [SER|_] = lists:filter(
diff --git a/test/tdf_SUITE.erl b/test/tdf_SUITE.erl
index d7d3ff9c..5915f051 100644
--- a/test/tdf_SUITE.erl
+++ b/test/tdf_SUITE.erl
@@ -1829,15 +1829,16 @@ tdf_app_id(Config) ->
 %%%===================================================================
 
 tdf_node_pid() ->
-    [[Pid]] = ets:match(gtp_context_reg, {{'cp-socket',{teid,'gtp-u','_'}},{ergw_sx_node, '$1'}}),
+    TEIDMatch = #socket_teid_key{name = 'cp-socket', type = 'gtp-u', _ = '_'},
+    [[Pid]] = ets:match(gtp_context_reg, {TEIDMatch, {ergw_sx_node, '$1'}}),
     Pid.
 
 tdf_seid() ->
-    [[SEID]] = ets:match(gtp_context_reg, {{seid, '$1'},{ergw_sx_node, '_'}}),
+    [[SEID]] = ets:match(gtp_context_reg, {#seid_key{seid = '$1'}, {ergw_sx_node, '_'}}),
     SEID.
 
 tdf_session_pid() ->
-    [[SEID]] = ets:match(gtp_context_reg, {{seid, '_'},{tdf, '$1'}}),
+    [[SEID]] = ets:match(gtp_context_reg, {#seid_key{_ = '_'}, {tdf, '$1'}}),
     SEID.
 
 ue_ip_address(Type, Config) ->