feat(erlang): make worker xdb version-aware for IPv4/IPv6

2026-06-28 00:01:09 +09:00 · 2026-06-28 00:01:09 +09:00 · 241831a6f0
parent 2f0dc5f33f
commit 241831a6f0
3 changed files with 156 additions and 61 deletions
--- 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),
--- 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(<<SPtr:32/little, EPtr:32/little, VectorIndexBin/binary>>, 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(<<SPtr:32/little, EPtr:32/little, VectorIndexBin/binary>>, 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, <<IpInt:32>>, 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;
        _ ->
-            <<A:8, B:8, _Rest/binary>> = <<IntIp:32>>,
+            <<A:8, B:8, _/binary>> = 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) ->
+    <<InputInt:32>> = 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, <<SIp:32/little, EIp:32/little, DataLen:16/little, DataPtr:32/little>>} = 
-	            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.
+    end.
+
+decode_segment_index(Bin, ?XDB_SEGMENT_INDEX_SIZE_V4) ->
+    <<SIp:32/little, EIp:32/little, DataLen:16/little, DataPtr:32/little>> = Bin,
+    {SIp, EIp, DataLen, DataPtr};
+decode_segment_index(Bin, ?XDB_SEGMENT_INDEX_SIZE_V6) ->
+    <<SIp:16/binary, EIp:16/binary, DataLen:16/little, DataPtr:32/little>> = 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.
--- a/binding/erlang/test/ip2region_worker_test.erl
+++ 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.