map_masq/imsi_list: bi-directional mapping, also change MO-ForwardSM-Arg

This introduces a second gb_tree of imsi mappings for the inverse
direction, in order to mangle MO-ForwardSM-Arg

2 files changed, 56 insertions, 17 deletions

diff --git a/src/imsi_list.erl b/src/imsi_list.erl
index 10dce97..005940b 100644
--- a/src/imsi_list.erl
+++ b/src/imsi_list.erl
@@ -1,7 +1,7 @@
 % Maintain a list of IMSIs in a gb_tree and match against it
 
-% (C) 2012 by Harald Welte <laforge@gnumonks.org>
-% (C) 2012 by On-Waves
+% (C) 2012-2013 by Harald Welte <laforge@gnumonks.org>
+% (C) 2012-2013 by On-Waves
 %
 % All Rights Reserved
 %
@@ -35,10 +35,14 @@
 -module(imsi_list).
 -author('Harald Welte <laforge@gnumonks.org>').
 
--export([read_file/1, read_list/1, match_imsi/2]).
+-export([read_file/1, read_list/1, match_imsi/2, match_imsi/3]).
+
+-record(state, {forward, reverse}).
 
 lines2tree(Iodev) ->
-	lines2tree(Iodev, gb_trees:empty()).
+	S = #state{forward = gb_trees:empty(),
+		   reverse = gb_trees:empty()},
+	lines2tree(Iodev, S).
 
 chomp(Line) when is_list(Line) ->
 	case lists:last(Line) of
@@ -50,12 +54,18 @@ chomp(Line) when is_list(Line) ->
 
 % 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].
+	[case string:to_integer([X]) of
+		{Int,[]} -> Int;
+		{error, F} ->
+			error_logger:error_report([{imsi_list_syntax_error,
+						Line2, {error, F}}]),
+			undefined
+	 end || X <- Line2].
 
-lines2tree(Iodev, Tree) ->
+lines2tree(Iodev, State) ->
 	case file:read_line(Iodev) of
 		eof ->
-			{ok, Tree};
+			{ok, State};
 		{error, Reason} ->
 			{error, Reason};
 		ebadf ->
@@ -67,7 +77,13 @@ lines2tree(Iodev, Tree) ->
 				[ImsiOldStr, ImsiNewStr] ->
 					ImsiOld = string_num_to_int_list(ImsiOldStr),
 					ImsiNew = string_num_to_int_list(ImsiNewStr),
-					lines2tree(Iodev, gb_trees:insert(ImsiOld, ImsiNew, Tree));
+					FwNew = gb_trees:insert(ImsiOld, ImsiNew,
+								State#state.forward),
+					RevNew = gb_trees:insert(ImsiNew, ImsiOld,
+								 State#state.reverse),
+					lines2tree(Iodev, #state{forward = FwNew,
+								 reverse = RevNew});
+				% FIXME: handle empty lines or skip bad lines
 				_ ->
 					{error, file_format}
 			end
@@ -84,15 +100,29 @@ read_file(FileName) ->
 	end.
 
 read_list(List) when is_list(List) ->
-	read_list(List, gb_trees:empty()).
+	S = #state{forward = gb_trees:empty(),
+		   reverse = gb_trees:empty()},
+	read_list(List, S).
 
 read_list([], Tree) ->
 	Tree;
-read_list([{Old, New}|Tail], Tree) ->
-	read_list(Tail, gb_trees:enter(Old, New, Tree)).
+read_list([{Old, New}|Tail], State) ->
+	FwNew = gb_trees:insert(Old, New, State#state.forward),
+	RevNew = gb_trees:insert(New, Old, State#state.reverse),
+	read_list(Tail, #state{forward = FwNew, reverse = RevNew}).
+
+match_imsi(State, Imsi) when is_list(Imsi) ->
+	match_imsi(forward, State, Imsi).
 
-match_imsi(Tree, Imsi) when is_list(Imsi) ->
-	case gb_trees:lookup(Imsi, Tree) of
+match_imsi(forward, State, Imsi) when is_list(Imsi) ->
+	case gb_trees:lookup(Imsi, State#state.forward) of
+		{value, ImsiNew} ->
+			{ok, ImsiNew};
+		none ->
+			{error, no_entry}
+	end;
+match_imsi(reverse, State, Imsi) when is_list(Imsi) ->
+	case gb_trees:lookup(Imsi, State#state.reverse) of
 		{value, ImsiNew} ->
 			{ok, ImsiNew};
 		none ->
diff --git a/src/map_masq.erl b/src/map_masq.erl
index 8ab4331..907d1c4 100644
--- a/src/map_masq.erl
+++ b/src/map_masq.erl
@@ -1,7 +1,7 @@
 % MAP masquerading application
 
-% (C) 2010-2012 by Harald Welte <laforge@gnumonks.org>
-% (C) 2010-2012 by On-Waves
+% (C) 2010-2013 by Harald Welte <laforge@gnumonks.org>
+% (C) 2010-2013 by On-Waves
 %
 % All Rights Reserved
 %
@@ -143,6 +143,10 @@ patch(From, #'LocationInfoWithLMSI'{'networkNode-Number' = NetNodeNr} = P) ->
 	NetNodeNrOut = patch_map_isdn_addr(From, NetNodeNr, msc),
 	P#'LocationInfoWithLMSI'{'networkNode-Number' = NetNodeNrOut};
 
+% MO-ForwardSM-Arg with optional IMSI
+patch(From, #'MO-ForwardSM-Arg'{imsi = ImsiIn} = P) ->
+	#'MO-ForwardSM-Arg'{imsi = patch_imsi(mo_fw_sm_arg, From, ImsiIn)};
+
 % patch the roaming number as it is sent from HLR to G-MSC (SRI Resp)
 patch(_From, {roamingNumber, RoamNumTBCD}) ->
 	RoamNumIn = map_codec:parse_addr_string(RoamNumTBCD),
@@ -504,15 +508,20 @@ generate_rewrite_entry({Name, MscSideInt, StpSideInt}) ->
 	StpSideList = osmo_util:int2digit_list(StpSideInt),
 	{Name, MscSideInt, StpSideInt, MscSideList, StpSideList}.
 
+imsi_direction(sri_sm_res) ->
+	forward;
+imsi_direction(mo_fw_sm_arg) ->
+	reverse.
 
 % check if we need to rewrite the IMSI
-patch_imsi(sri_sm_res, from_msc, ImsiIn) ->
+patch_imsi(MsgType, from_msc, ImsiIn) ->
+	IsForward = imsi_direction(MsgType),
 	case application:get_env(mgw_nat, imsi_rewrite_tree) of
 		{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
+			case imsi_list:match_imsi(IsForward, ImsiTree, Imsi) of
 				{ok, NewImsi} ->
 					map_codec:encode_map_tbcd(NewImsi);
 				_ ->