diff --git a/binding/erlang/src/ip2region_sup.erl b/binding/erlang/src/ip2region_sup.erl index 304bc98..6aa6695 100644 --- a/binding/erlang/src/ip2region_sup.erl +++ b/binding/erlang/src/ip2region_sup.erl @@ -42,9 +42,23 @@ init([]) -> %% create_table() -> Opts = [named_table, set, public, {read_concurrency, true}, {keypos, 1}], - ets:new(?XDB_VECTOR_INDEX, Opts), - ets:new(?XDB_SEGMENT_INDEX, Opts), - ets:new(?IP2REGION_CACHE, Opts). + %% Legacy tables (kept for backward compatibility) + ensure_table(?XDB_VECTOR_INDEX, Opts), + ensure_table(?XDB_SEGMENT_INDEX, Opts), + ensure_table(?IP2REGION_CACHE, Opts), + %% Version-specific tables for dual-stack support + ensure_table(?XDB_VECTOR_INDEX_V4, Opts), + ensure_table(?XDB_VECTOR_INDEX_V6, Opts), + ensure_table(?XDB_SEGMENT_INDEX_V4, Opts), + ensure_table(?XDB_SEGMENT_INDEX_V6, Opts), + ensure_table(?IP2REGION_CACHE_V4, Opts), + ensure_table(?IP2REGION_CACHE_V6, Opts). + +ensure_table(Name, Opts) -> + case ets:whereis(Name) of + undefined -> ets:new(Name, Opts); + _ -> ok + end. start_ip2region_pool(Sup) -> {ok, PoolArgsCfg} = application:get_env(poolargs), diff --git a/binding/erlang/src/ip2region_worker.erl b/binding/erlang/src/ip2region_worker.erl index 64d039c..e207635 100644 --- a/binding/erlang/src/ip2region_worker.erl +++ b/binding/erlang/src/ip2region_worker.erl @@ -2,9 +2,9 @@ %% Copyright 2022 The Ip2Region Authors. All rights reserved. %% Use of this source code is governed by a Apache2.0-style %% license that can be found in the LICENSE file. -%% -%% @doc -%% ip2region xdb client worker +%% +%% @doc +%% ip2region xdb client worker, now version-aware (IPv4/IPv6). %% @end %%%------------------------------------------------------------------- -module(ip2region_worker). @@ -15,7 +15,12 @@ -export([start/1, stop/1, start_link/1]). -export([search/2]). -export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2, code_change/3]). --record(state, {xdb_fd}). + +-record(state, { + xdb_fd, + version :: ipv4 | ipv6, + segment_index_size :: pos_integer() +}). %%========================================== %% API @@ -28,7 +33,6 @@ start_link(Args) -> Opts = [{spawn_opt, [{min_heap_size, 6000}]}], gen_server:start_link(?MODULE, Args, Opts). - stop(Pid) -> gen_server:call(Pid, stop). @@ -38,26 +42,42 @@ search(Pid, Ip) -> %%========================================== %% gen_server callbacks %% ========================================= -init(_Args) -> +init(Args) -> process_flag(trap_exit, true), - AppName = + AppName = case application:get_application() of {ok, AName} -> AName; _ -> ?APP_NAME end, PrivDir = code:priv_dir(AppName), - XdbFileName = filename:join([PrivDir, "ip2region.xdb"]), + XdbFileName = + case proplists:get_value(xdb_file, Args) of + undefined -> filename:join([PrivDir, "ip2region.xdb"]); + Path -> + case filename:pathtype(Path) of + absolute -> Path; + _ -> filename:join([PrivDir, Path]) + end + end, error_logger:info_report(io_lib:format("XdbFile:~s~n", [XdbFileName])), {ok, IoDevice} = file:open(XdbFileName, [read, binary]), - load_vector_index(IoDevice), - {ok, #state{xdb_fd = IoDevice}}. + {ok, HeaderBin} = file:read(IoDevice, ?XDB_HEADER_SIZE), + {ok, Header} = ip2region_xdb:parse_header(HeaderBin), + Version = resolve_version(Header), + SegmentIndexSize = ip2region_xdb:segment_index_size(Version), + load_vector_index(IoDevice, Version), + {ok, #state{ + xdb_fd = IoDevice, + version = Version, + segment_index_size = SegmentIndexSize + }}. handle_call(Request, From, State) -> try do_call(Request, From, State) catch Class:Error:Stacktrace -> - error_logger:error_report(io_lib:format("~p handle call error, Req:~p ~p, stacktrace:~p~n", + error_logger:error_report(io_lib:format("~p handle call error, Req:~p ~p, stacktrace:~p~n", [?MODULE, Request, {Class, Error}, Stacktrace])), {reply, {error, {Class, Error}}, State} end. @@ -67,7 +87,7 @@ handle_cast(Msg, State) -> do_cast(Msg, State) catch Class:Error:Stacktrace -> - error_logger:error_report(io_lib:format("~p handle cast error, Msg:~p, ~p, stacktrace:~w~n", + error_logger:error_report(io_lib:format("~p handle cast error, Msg:~p, ~p, stacktrace:~w~n", [?MODULE, Msg, {Class, Error}, Stacktrace])), {noreply, State} end. @@ -77,7 +97,7 @@ handle_info(Info, State) -> do_info(Info, State) catch Class:Error:Stacktrace -> - error_logger:error_report(io_lib:format("~p handle info error, Info:~p, ~p, stacktrace:~p~n", + error_logger:error_report(io_lib:format("~p handle info error, Info:~p, ~p, stacktrace:~p~n", [?MODULE, Info, {Class, Error}, Stacktrace])), {noreply, State} end. @@ -85,22 +105,19 @@ handle_info(Info, State) -> terminate(_Reason, State) -> #state{xdb_fd = XdbFd} = State, case is_pid(XdbFd) of - true -> - file:close(XdbFd); - _ -> - skip + true -> file:close(XdbFd); + _ -> skip end, ok. code_change(_OldVsn, State, _Extra) -> {ok, State}. - %%========================================== %% Internal function %% ========================================= do_call({search, Ip}, _From, #state{xdb_fd = IoDevice} = State) -> - Reply = search_ip(IoDevice, Ip), + Reply = search_ip(IoDevice, Ip, State), {reply, Reply, State}; do_call(stop, _From, State) -> @@ -117,67 +134,118 @@ do_cast(Msg, State) -> do_info(Info, State) -> error_logger:error_report(io:format("unknown info: ~p", [Info])), {noreply, State}. - -load_vector_index(IoDevice) -> - Key = ip2region_header_loaded, - case persistent_term:get(Key, false) of - true -> ok; - _ -> - {ok, <<_Header:?XDB_HEADER_SIZE/binary, VectorIndexBin/binary>> } = - file:read(IoDevice, ?XDB_HEADER_SIZE + ?XDB_VECTOR_INDEX_COUNT*8), - load_vector_index_aux(VectorIndexBin, 0), - persistent_term:put(Key, true) +resolve_version(Header) -> + case ip2region_xdb:header_version(Header) of + 2 -> ipv4; + 3 -> + case ip2region_xdb:header_ip_version(Header) of + ?IP_VERSION_4 -> ipv4; + ?IP_VERSION_6 -> ipv6; + _ -> ipv4 + end; + _ -> ipv4 end. -load_vector_index_aux(<<>>, _Index) -> ok; -load_vector_index_aux(<>, Index) -> - Term = {Index, SPtr, EPtr}, - ets:insert(?XDB_VECTOR_INDEX, Term), - load_vector_index_aux(VectorIndexBin, Index + 1). +load_vector_index(IoDevice, Version) -> + Table = vector_index_table(Version), + case ets:info(Table, size) of + undefined -> + Opts = [named_table, set, public, {read_concurrency, true}, {keypos, 1}], + ets:new(Table, Opts), + load_vector_index_data(IoDevice, Table); + 0 -> + load_vector_index_data(IoDevice, Table); + _ -> + ok + end. +load_vector_index_data(IoDevice, Table) -> + {ok, VectorIndexBin} = + file:read(IoDevice, ?XDB_VECTOR_INDEX_COUNT * 8), + load_vector_index_aux(VectorIndexBin, 0, Table). -search_ip(IoDevice, Ip) -> - IntIp = ip2region_util:ipv4_to_n(Ip), - case ets:lookup(?IP2REGION_CACHE, IntIp) of - [{_IntIp, RegionInfo}] -> +load_vector_index_aux(<<>>, _Index, _Table) -> ok; +load_vector_index_aux(<>, Index, Table) -> + ets:insert(Table, {Index, SPtr, EPtr}), + load_vector_index_aux(VectorIndexBin, Index + 1, Table). + +search_ip(IoDevice, IpInt, State) when is_integer(IpInt) -> + search_ip(IoDevice, <>, State); +search_ip(IoDevice, Ip, #state{version = Version, segment_index_size = SegSize}) -> + CacheTable = cache_table(Version), + VectorTable = vector_index_table(Version), + SegmentTable = segment_index_table(Version), + case ets:lookup(CacheTable, Ip) of + [{_, RegionInfo}] -> RegionInfo; _ -> - <> = <>, + <> = Ip, VectorIdx = A * ?XDB_VECTOR_COLS + B, - [{_, SPtr, EPtr}] = ets:lookup(?XDB_VECTOR_INDEX, VectorIdx), - RegionInfo = search_ip(IoDevice, IntIp, SPtr, EPtr, 0, (EPtr - SPtr) div ?XDB_SEGMENT_INDEX_SIZE), - ets:insert_new(?IP2REGION_CACHE, {IntIp, RegionInfo}), + [{_, SPtr, EPtr}] = ets:lookup(VectorTable, VectorIdx), + RegionInfo = search_ip(IoDevice, Ip, SPtr, EPtr, 0, + (EPtr - SPtr) div SegSize, SegSize, Version, SegmentTable), + ets:insert_new(CacheTable, {Ip, RegionInfo}), RegionInfo end. -search_ip(IoDevice, IntIp, SPtr, EPtr, Low, High) when Low =< High -> +search_ip(IoDevice, Ip, SPtr, EPtr, Low, High, SegSize, Version, SegmentTable) when Low =< High -> Middle = (Low + High) bsr 1, - SPtr2 = SPtr + Middle * ?XDB_SEGMENT_INDEX_SIZE, - {SIp, EIp, DataLen, DataPtr} = read_segement_index(IoDevice, SPtr2), - if - IntIp < SIp -> - search_ip(IoDevice, IntIp, SPtr, EPtr, Low, Middle - 1); - IntIp > EIp -> - search_ip(IoDevice, IntIp, SPtr, EPtr, Middle + 1, High); - true -> + SPtr2 = SPtr + Middle * SegSize, + {SIp, EIp, DataLen, DataPtr} = read_segment_index(IoDevice, SPtr2, SegSize, SegmentTable), + case ip_in_range(Ip, SIp, EIp, Version) of + below -> + search_ip(IoDevice, Ip, SPtr, EPtr, Low, Middle - 1, SegSize, Version, SegmentTable); + above -> + search_ip(IoDevice, Ip, SPtr, EPtr, Middle + 1, High, SegSize, Version, SegmentTable); + inside -> {ok, DataBin} = read_file(IoDevice, DataPtr, DataLen), unicode:characters_to_nfc_list(DataBin) end; -search_ip(_IoDevice, _IntIp, _SPtr, _EPtr, _Low, _High) -> +search_ip(_IoDevice, _Ip, _SPtr, _EPtr, _Low, _High, _SegSize, _Version, _SegmentTable) -> {error, unknown}. +ip_in_range(Ip, SIp, EIp, ipv4) -> + <> = Ip, + if + InputInt < SIp -> below; + InputInt > EIp -> above; + true -> inside + end; +ip_in_range(Ip, SIp, EIp, ipv6) -> + if + Ip < SIp -> below; + Ip > EIp -> above; + true -> inside + end. + read_file(IoDevice, Position, DataLength) -> file:position(IoDevice, {bof, Position}), file:read(IoDevice, DataLength). -read_segement_index(IoDevice, SPtr) -> - case ets:lookup(?XDB_SEGMENT_INDEX, SPtr) of +read_segment_index(IoDevice, SPtr, SegSize, SegmentTable) -> + case ets:lookup(SegmentTable, SPtr) of [{_SPtr, SIp, EIp, DataLen, DataPtr}] -> {SIp, EIp, DataLen, DataPtr}; _ -> - {ok, <>} = - read_file(IoDevice, SPtr, ?XDB_SEGMENT_INDEX_SIZE), - ets:insert_new(?XDB_SEGMENT_INDEX, {SPtr, SIp, EIp, DataLen, DataPtr}), + {ok, Bin} = read_file(IoDevice, SPtr, SegSize), + {SIp, EIp, DataLen, DataPtr} = decode_segment_index(Bin, SegSize), + ets:insert_new(SegmentTable, {SPtr, SIp, EIp, DataLen, DataPtr}), {SIp, EIp, DataLen, DataPtr} - end. \ No newline at end of file + end. + +decode_segment_index(Bin, ?XDB_SEGMENT_INDEX_SIZE_V4) -> + <> = Bin, + {SIp, EIp, DataLen, DataPtr}; +decode_segment_index(Bin, ?XDB_SEGMENT_INDEX_SIZE_V6) -> + <> = Bin, + {SIp, EIp, DataLen, DataPtr}. + +vector_index_table(ipv4) -> ?XDB_VECTOR_INDEX_V4; +vector_index_table(ipv6) -> ?XDB_VECTOR_INDEX_V6. + +segment_index_table(ipv4) -> ?XDB_SEGMENT_INDEX_V4; +segment_index_table(ipv6) -> ?XDB_SEGMENT_INDEX_V6. + +cache_table(ipv4) -> ?IP2REGION_CACHE_V4; +cache_table(ipv6) -> ?IP2REGION_CACHE_V6. diff --git a/binding/erlang/test/ip2region_worker_test.erl b/binding/erlang/test/ip2region_worker_test.erl new file mode 100644 index 0000000..5a2417c --- /dev/null +++ b/binding/erlang/test/ip2region_worker_test.erl @@ -0,0 +1,13 @@ +-module(ip2region_worker_test). +-include_lib("eunit/include/eunit.hrl"). +-include("ip2region.hrl"). + +worker_search_v4_binary_test() -> + ip2region_sup:create_table(), + {ok, Pid} = ip2region_worker:start_link([{xdb_file, "ip2region.xdb"}]), + try + Region = ip2region_worker:search(Pid, <<1,0,8,0>>), + ?assert(is_list(Region)) + after + ip2region_worker:stop(Pid) + end.