feat(erlang): auto-route search by IPv4/IPv6 to correct worker pool

binding/erlang/src/xdb.erl

@@ -1,11 +1,10 @@
-
 %%%-------------------------------------------------------------------
 %% 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 search api
+%% ip2region xdb client search api with IPv4/IPv6 auto-routing.
 %% @end
 %%%-------------------------------------------------------------------
 -module(xdb).
@@ -13,20 +12,36 @@
 
 -export([search/1]).
 
--spec search(Ip :: tuple() | list() | binary()) -> Result :: binary | {error, Reason::atom()}.
+-spec search(Ip :: tuple() | list() | binary() | integer()) ->
-search(Ip) when is_integer(Ip); is_list(Ip); is_tuple(Ip); is_binary(Ip) ->
+    Result :: list() | {error, Reason::atom()}.
-    case ip2region_util:ipv4_to_n(Ip) of
+search(Ip) ->
-    IntIp when is_integer(IntIp) ->
+    case ip2region_util:ip_to_bytes(Ip) of
-        case ets:lookup(?IP2REGION_CACHE, IntIp) of
+        {ok, ipv4, IpBin} ->
-        [{_IntIp, Region}] -> Region;
+            do_search(v4_pool(), ipv4, IpBin);
-        _ ->
+        {ok, ipv6, IpBin} ->
-            Worker = poolboy:checkout(?IP2REGION_POOL, true, infinity),
+            do_search(?IP2REGION_POOL_V6, ipv6, IpBin);
-            try
-                ip2region_worker:search(Worker, IntIp)
-            after
-                poolboy:checkin(?IP2REGION_POOL, Worker)
-            end
-        end;
         Ret ->
             Ret
     end.
+
+do_search(PoolName, Version, IpBin) ->
+    CacheTable = cache_table(Version),
+    case ets:lookup(CacheTable, IpBin) of
+        [{_, Region}] -> Region;
+        _ ->
+            Worker = poolboy:checkout(PoolName, true, infinity),
+            try
+                ip2region_worker:search(Worker, IpBin)
+            after
+                poolboy:checkin(PoolName, Worker)
+            end
+    end.
+
+v4_pool() ->
+    case whereis(?IP2REGION_POOL) of
+        Pid when is_pid(Pid) -> ?IP2REGION_POOL;
+        undefined -> ?IP2REGION_POOL_V4
+    end.
+
+cache_table(ipv4) -> ?IP2REGION_CACHE_V4;
+cache_table(ipv6) -> ?IP2REGION_CACHE_V6.

binding/erlang/test/xdb_test.erl

@@ -4,7 +4,7 @@
 
 search_test_() ->
     application:ensure_started(ip2region),
-    A = "中国|0|广东省|广州市|电信",
+    A = "中国|广东省|广州市|中国电信|CN",
     Region0 = xdb:search("1.0.8.0"),
     Region1 = xdb:search(<<"1.0.8.0">>),
     Region2 = xdb:search({1,0,8,0}),
@@ -14,5 +14,20 @@ search_test_() ->
         ?_assert(A =:= Region1),
         ?_assert(A =:= Region2),
         ?_assert({error, bad_ip_format} =:= Region3)
-
+    ].
+
+ipv6_search_test_() ->
+    application:ensure_started(ip2region),
+    [
+        ?_assert(is_list(xdb:search("2001:4860:4860::8888"))),
+        ?_assert(is_list(xdb:search(<<"2001:4860:4860::8888">>))),
+        ?_assert(is_list(xdb:search({8193, 10304, 10304, 0, 0, 0, 0, 34952})))
+    ].
+
+invalid_search_test_() ->
+    application:ensure_started(ip2region),
+    [
+        ?_assertEqual({error, bad_ip_format}, xdb:search("xxx.0.8.0")),
+        ?_assertEqual({error, bad_ip_format}, xdb:search("::ggg")),
+        ?_assertEqual({error, bad_ip_format}, xdb:search({1,2,3}))
     ].