From fce6cf18ddb85a6b5f1ffd73636b995bd74bcf60 Mon Sep 17 00:00:00 2001
From: Harald Welte <laforge@gnumonks.org>
Date: Mon, 20 Feb 2012 07:56:16 +0100
Subject: [PATCH] IMSI rewrite: change from prefix rewrite to full N:M
 rewriting

---
 src/imsi_list.erl       | 26 +++++++++++++++++---------
 src/map_masq.erl        | 24 ++++++------------------
 test/map_masq_tests.erl | 15 ++++++++-------
 3 files changed, 31 insertions(+), 34 deletions(-)

diff --git a/src/imsi_list.erl b/src/imsi_list.erl
index 2f60673..8fe1827 100644
--- a/src/imsi_list.erl
+++ b/src/imsi_list.erl
@@ -35,6 +35,10 @@ chomp(Line) when is_list(Line) ->
 			Line
 	end.
 
+% convert from "12345" to [1,2,3,4,5]
+string_num_to_int_list(Line2) ->
+	[case string:to_integer([X]) of {Int,[]} -> Int end || X <- Line2].
+
 lines2tree(Iodev, Tree) ->
 	case file:read_line(Iodev) of
 		eof ->
@@ -46,8 +50,14 @@ lines2tree(Iodev, Tree) ->
 		{ok, Line} ->
 			% FIXME: convert to digit list
 			Line2 = chomp(Line),
-			Line3 = [case string:to_integer([X]) of {Int,[]} -> Int end || X <- Line2],
-			lines2tree(Iodev, gb_trees:insert(Line3, true, Tree))
+			case string:tokens(Line2, ",;") of
+				[ImsiOldStr, ImsiNewStr] ->
+					ImsiOld = string_num_to_int_list(ImsiOldStr),
+					ImsiNew = string_num_to_int_list(ImsiNewStr),
+					lines2tree(Iodev, gb_trees:insert(ImsiOld, ImsiNew, Tree));
+				_ ->
+					{error, file_format}
+			end
 	end.
 
 
@@ -65,15 +75,13 @@ read_list(List) when is_list(List) ->
 
 read_list([], Tree) ->
 	Tree;
-read_list([Head|Tail], Tree) ->
-	read_list(Tail, gb_trees:enter(Head, true, Tree)).
+read_list([{Old, New}|Tail], Tree) ->
+	read_list(Tail, gb_trees:enter(Old, New, Tree)).
 
 match_imsi(Tree, Imsi) when is_list(Imsi) ->
 	case gb_trees:lookup(Imsi, Tree) of
-		{value, true} ->
-			true;
-		{value, _} ->
-			false;
+		{value, ImsiNew} ->
+			{ok, ImsiNew};
 		none ->
-			false
+			{error, no_entry}
 	end.
diff --git a/src/map_masq.erl b/src/map_masq.erl
index 410243d..29346fc 100644
--- a/src/map_masq.erl
+++ b/src/map_masq.erl
@@ -462,11 +462,9 @@ config_update() ->
 	case application:get_env(mgw_nat, imsi_rewrite_file) of
 		{ok, ImsiListFile} ->
 			{ok, ImsiTree} = imsi_list:read_file(ImsiListFile),
-			{ok, OldPrefix} = application:get_env(mgw_nat, imsi_rewrite_old_prefix),
-			{ok, NewPrefix} = application:get_env(mgw_nat, imsi_rewrite_new_prefix),
-			io:format("(Re)generated IMSI rewrite table: ~p entries, ~p -> ~p~n",
-				  [gb_trees:size(ImsiTree), OldPrefix, NewPrefix]),
-			application:set_env(mgw_nat, imsi_rewrite_tree, {ImsiTree, OldPrefix, NewPrefix});
+			io:format("(Re)generated IMSI rewrite table: ~p entries~n",
+				  [gb_trees:size(ImsiTree)]),
+			application:set_env(mgw_nat, imsi_rewrite_tree, ImsiTree);
 		_ ->
 			ok
 	end,
@@ -498,15 +496,14 @@ generate_rewrite_entry({Name, MscSideInt, StpSideInt}) ->
 % check if we need to rewrite the IMSI
 patch_imsi(sri_sm_res, from_msc, ImsiIn) ->
 	case application:get_env(mgw_nat, imsi_rewrite_tree) of
-		{ok, {ImsiTree, OldPrefix, NewPrefix}} ->
+		{ok, ImsiTree} ->
 			% decode IMSI into list of digits
 			Imsi = map_codec:parse_map_addr(ImsiIn),
 			% rewrite prefix, if it matches
 			case imsi_list:match_imsi(ImsiTree, Imsi) of
-				true ->
-					NewImsi = prefix_rewrite(OldPrefix, NewPrefix, Imsi),
+				{ok, NewImsi} ->
 					map_codec:encode_map_tbcd(NewImsi);
-				false ->
+				_ ->
 					ImsiIn
 			end;
 		_ ->
@@ -514,12 +511,3 @@ patch_imsi(sri_sm_res, from_msc, ImsiIn) ->
 	end;
 patch_imsi(_, _, Imsi) ->
 	Imsi.
-
-prefix_rewrite(OldPrefix, NewPrefix, Imsi) when is_list(OldPrefix), is_list(NewPrefix), is_list(Imsi) ->
-	case lists:sublist(Imsi, length(OldPrefix)) of
-		OldPrefix ->
-			% remove old prefix and prepend new prefix
-			NewPrefix ++ lists:sublist(Imsi, length(OldPrefix)+1, length(Imsi));
-		_ ->
-			Imsi
-	end.
diff --git a/test/map_masq_tests.erl b/test/map_masq_tests.erl
index 1f37d97..5943741 100644
--- a/test/map_masq_tests.erl
+++ b/test/map_masq_tests.erl
@@ -6,19 +6,20 @@
 -define(OLD_PFX, [2,6,2,7,7]).
 -define(NEW_PFX, [9,0,1,7,7]).
 
--define(IMSI_TRUE, [0,0,0,0,0,0,0,0,0,9]).
--define(IMSI_FALSE, [1,0,0,0,0,0,0,0,0,1]).
--define(SRI_SM_MATCH_IN, gen_sri_sm(?OLD_PFX ++ ?IMSI_TRUE)).
--define(SRI_SM_MATCH_OUT, gen_sri_sm(?NEW_PFX ++ ?IMSI_TRUE)).
--define(SRI_SM_NOMATCH_IN, gen_sri_sm(?OLD_PFX ++ ?IMSI_FALSE)).
+-define(IMSI_TRUE, [2,6,2,7,7,0,0,0,0,0,0,0,0,0,9]).
+-define(IMSI_NEW, [9,0,1,7,7,0,0,0,0,0,0,0,0,0,9]).
+-define(IMSI_FALSE, [2,6,2,6,6,1,0,0,0,0,0,0,0,0,1]).
+-define(SRI_SM_MATCH_IN, gen_sri_sm(?IMSI_TRUE)).
+-define(SRI_SM_MATCH_OUT, gen_sri_sm(?IMSI_NEW)).
+-define(SRI_SM_NOMATCH_IN, gen_sri_sm(?IMSI_FALSE)).
 
 gen_sri_sm(Imsi) when is_list(Imsi) ->
 	{'begin', #'MapSpecificPDUs_begin'{otid=[0,0,0,1], components=[{basicROS, {returnResult, #'MapSpecificPDUs_end_components_SEQOF_basicROS_returnResult'{invokeId = 1, result=#'MapSpecificPDUs_end_components_SEQOF_basicROS_returnResult_result'{opcode={local, 2342}, result=#'RoutingInfoForSM-Res'{imsi=map_codec:encode_map_tbcd(Imsi), _ = asn1_NOVALUE}, _ = asn1_NOVALUE}, _ = asn1_NOVALUE}}}], _ = asn1_NOVALUE}}.
 
 
 setup() ->
-	Tree = imsi_list:read_list([?OLD_PFX ++ ?IMSI_TRUE]),
-	application:set_env(mgw_nat, imsi_rewrite_tree, {Tree, ?OLD_PFX, ?NEW_PFX}).
+	Tree = imsi_list:read_list([{?IMSI_TRUE, ?IMSI_NEW}]),
+	application:set_env(mgw_nat, imsi_rewrite_tree, Tree).
 
 teardown(_) ->
 	application:unset_env(mgw_nat, imsi_rewrite_tree).