diff --git a/README.md b/README.md index 81a18aa..fc15a50 100644 --- a/README.md +++ b/README.md @@ -318,10 +318,10 @@ If list is empty, no limits will be checked. Following policies are supported: -* `{max_connections, KEYS, NUMBER}` - EXPERIMENTAL! if there are more than NUMBER connections with - KEYS to the proxy, new connections with those KEYS will be rejected * `{in_table, KEY, TABLE_NAME}` - only allow connections if KEY is present in TABLE_NAME (whitelist) * `{not_in_table, KEY, TABLE_NAME}` - only allow connections if KEY is *not* present in TABLE_NAME (blacklist) +* `{max_connections, KEYS, NUMBER}` - EXPERIMENTAL! if there are more than NUMBER connections with + KEYS to the proxy, new connections with those KEYS will be rejected. Where: diff --git a/src/mtp_handler.erl b/src/mtp_handler.erl index dc418fa..c4345c7 100644 --- a/src/mtp_handler.erl +++ b/src/mtp_handler.erl @@ -214,12 +214,18 @@ handle_info(Other, S) -> {noreply, S}. terminate(_Reason, #state{started_at = Started, listener = Listener, - addr = {Ip, _}, policy_state = MaybeSni} = S) -> - try mtp_policy:dec( - application:get_env(?APP, policy, []), - Listener, Ip, MaybeSni) - catch T:R -> - ?log(warning, "Failed to decrement policy: ~p:~p", [T, R]) + addr = {Ip, _}, policy_state = PolicyState} = S) -> + case PolicyState of + {ok, TlsDomain} -> + try mtp_policy:dec( + application:get_env(?APP, policy, []), + Listener, Ip, TlsDomain) + catch T:R -> + ?log(warning, "Failed to decrement policy: ~p:~p", [T, R]) + end; + _ -> + %% Failed before policy was stored in state. Eg, because of "policy_error" + ok end, maybe_close_down(S), mtp_metric:count_inc([?APP, in_connection_closed, total], 1, #{labels => [Listener]}), @@ -306,25 +312,26 @@ parse_upstream_data(<> = AllData, check_tls_policy(Listener, Ip, Meta), Codec1 = mtp_codec:replace(tls, true, TlsCodec, Codec0), Codec = mtp_codec:push_back(tls, Tail, Codec1), - ok = up_send_raw(Response, S), + ok = up_send_raw(Response, S), %FIXME: if this send fail, we will get counter policy leak {ok, S#state{codec = Codec, stage = init, - policy_state = maps:get(sni_domain, Meta, undefined)}}; + policy_state = {ok, maps:get(sni_domain, Meta, undefined)}}}; parse_upstream_data(<> = Data, #state{stage = init} = S) -> parse_upstream_data(Data, S#state{stage = tls_hello}); parse_upstream_data(<>, #state{stage = init, secret = Secret, listener = Listener, codec = Codec0, - ad_tag = Tag, addr = {Ip, _} = Addr} = S) -> + ad_tag = Tag, addr = {Ip, _} = Addr, policy_state = PState0} = S) -> case mtp_obfuscated:from_header(Header, Secret) of {ok, DcId, PacketLayerMod, CryptoCodecSt} -> maybe_check_replay(Header), - ProtoToReport = + {ProtoToReport, PState} = case mtp_codec:info(tls, Codec0) of {true, _} when PacketLayerMod == mtp_secure -> - mtp_secure_fake_tls; + {mtp_secure_fake_tls, PState0}; {false, _} -> assert_protocol(PacketLayerMod), check_policy(Listener, Ip, undefined), - PacketLayerMod + %FIXME: if any codebelow fail, we will get counter policy leak + {PacketLayerMod, {ok, undefined}} end, mtp_metric:count_inc([?APP, protocol_ok, total], 1, #{labels => [Listener, ProtoToReport]}), @@ -341,6 +348,7 @@ parse_upstream_data(<>, S#state{down = Downstream, dc_id = {RealDcId, Pool}, codec = Codec, + policy_state = PState, stage = tunnel}, hibernate)); {error, Reason} when is_atom(Reason) -> diff --git a/src/mtp_policy.erl b/src/mtp_policy.erl index 2889cf7..df5e5fc 100644 --- a/src/mtp_policy.erl +++ b/src/mtp_policy.erl @@ -49,13 +49,12 @@ check(Rules, ListenerName, ClientIp, TlsDomain) -> end, Rules). dec(Rules, ListenerName, ClientIp,TlsDomain) -> - %% FIXME: this is not idempotent if `check/4` returned `false`! Vars = vars(ListenerName, ClientIp,TlsDomain), lists:foreach( fun({max_connections, Keys, _Max}) -> try Key = [val(K, Vars) || K <- Keys], - mtp_policy_counter:increment(Key) + mtp_policy_counter:decrement(Key) catch throw:not_applicable -> ok end; @@ -77,6 +76,7 @@ check({max_connections, Keys, Max}, Vars) -> Key = [val(K, Vars) || K <- Keys], case mtp_policy_counter:increment(Key) of N when N > Max -> + mtp_policy_counter:decrement(Key), false; _ -> true diff --git a/src/mtp_policy_counter.erl b/src/mtp_policy_counter.erl index a4636b8..90ba3ad 100644 --- a/src/mtp_policy_counter.erl +++ b/src/mtp_policy_counter.erl @@ -15,6 +15,7 @@ -export([start_link/0]). -export([increment/1, decrement/1, + get/1, flush/0]). %% gen_server callbacks @@ -35,7 +36,7 @@ increment(Key) -> -spec decrement(key()) -> integer(). decrement(Key) -> - try ets:update_counter(?TAB, Key, 1) of + try ets:update_counter(?TAB, Key, -1) of New when New =< 0 -> ets:delete(?TAB, Key), 0; @@ -45,6 +46,13 @@ decrement(Key) -> 0 end. +-spec get(key()) -> non_neg_integer(). +get(Key) -> + case ets:lookup(?TAB, Key) of + [] -> 0; + [{_, V}] -> V + end. + %% @doc Clean all counters flush() -> gen_server:call(?MODULE, flush). diff --git a/test/single_dc_SUITE.erl b/test/single_dc_SUITE.erl index 7b33e7a..a164e6d 100644 --- a/test/single_dc_SUITE.erl +++ b/test/single_dc_SUITE.erl @@ -7,16 +7,18 @@ init_per_testcase/2, end_per_testcase/2]). --export([echo_secure_case/1, - echo_abridged_many_packets_case/1, - echo_tls_case/1, - packet_too_large_case/1, +-export([config_change_case/1, downstream_size_backpressure_case/1, downstream_qlen_backpressure_case/1, - config_change_case/1, + echo_secure_case/1, + echo_abridged_many_packets_case/1, + echo_tls_case/1, + ipv6_connect_case/1, + packet_too_large_case/1, + policy_max_conns_case/1, + policy_whitelist_case/1, replay_attack_case/1, - replay_attack_server_error_case/1, - ipv6_connect_case/1 + replay_attack_server_error_case/1 ]). -export([set_env/2, @@ -67,11 +69,11 @@ echo_secure_case(Cfg) when is_list(Cfg) -> Port = ?config(mtp_port, Cfg), Secret = ?config(mtp_secret, Cfg), Cli = mtp_test_client:connect(Host, Port, Secret, DcId, mtp_secure), - Data = crypto:strong_rand_bytes(64), - Cli1 = mtp_test_client:send(Data, Cli), - {ok, Packet, Cli2} = mtp_test_client:recv_packet(Cli1, 1000), + Cli2 = ping(Cli), + ?assertEqual( + 1, mtp_test_metric:get_tags( + count, [?APP, protocol_ok, total], [?FUNCTION_NAME, mtp_secure])), ok = mtp_test_client:close(Cli2), - ?assertEqual(Data, Packet), ok = mtp_test_metric:wait_for_value( count, [?APP, in_connection_closed, total], [?FUNCTION_NAME], 1, 5000), ?assertEqual(1, mtp_test_metric:get_tags( @@ -129,11 +131,11 @@ echo_tls_case(Cfg) when is_list(Cfg) -> Port = ?config(mtp_port, Cfg), Secret = ?config(mtp_secret, Cfg), Cli0 = mtp_test_client:connect(Host, Port, Secret, DcId, {mtp_fake_tls, <<"example.com">>}), - Data = crypto:strong_rand_bytes(64), - Cli1 = mtp_test_client:send(Data, Cli0), - {ok, Packet, Cli2} = mtp_test_client:recv_packet(Cli1, 1000), - ok = mtp_test_client:close(Cli2), - ?assertEqual(Data, Packet). + Cli1 = ping(Cli0), + ?assertEqual( + 1, mtp_test_metric:get_tags( + count, [?APP, protocol_ok, total], [?FUNCTION_NAME, mtp_secure_fake_tls])), + ok = mtp_test_client:close(Cli1). %% @doc test that client trying to send too big packets will be force-disconnected @@ -412,15 +414,88 @@ ipv6_connect_case(Cfg) when is_list(Cfg) -> ?assertEqual(not_found, ConnCount()), ?assertEqual(8, tuple_size(Host)), Cli0 = mtp_test_client:connect(Host, Port, Secret, DcId, mtp_secure), - Data = crypto:strong_rand_bytes(64), - Cli1 = mtp_test_client:send(Data, Cli0), - {ok, Packet, Cli2} = mtp_test_client:recv_packet(Cli1, 1000), - ok = mtp_test_client:close(Cli2), - ?assertEqual(Data, Packet), + Cli1 = ping(Cli0), + ok = mtp_test_client:close(Cli1), ?assertEqual(1, ConnCount()), ok = mtp_test_metric:wait_for_value( count, [?APP, in_connection_closed, total], [?FUNCTION_NAME], 1, 5000). + +%% @doc Test "max_connections" policy +policy_max_conns_case({pre, Cfg}) -> + Cfg1 = setup_single(?FUNCTION_NAME, 10000 + ?LINE, #{}, Cfg), + %% Allow max 2 connections from IP + set_env([{policy, [{max_connections, [port_name, client_ipv4], 2}]}], Cfg1); +policy_max_conns_case({post, Cfg}) -> + stop_single(Cfg), + reset_env(Cfg); +policy_max_conns_case(Cfg) when is_list(Cfg) -> + DcId = ?config(dc_id, Cfg), + Host = ?config(mtp_host, Cfg), + Port = ?config(mtp_port, Cfg), + Secret = ?config(mtp_secret, Cfg), + SureClose = + fun(Cli) -> + PreClosed = + mtp_test_metric:get_tags( + count, [?APP, in_connection_closed, total], [?FUNCTION_NAME]), + ok = mtp_test_client:close(Cli), + ok = mtp_test_metric:wait_for_value( + count, [?APP, in_connection_closed, total], [?FUNCTION_NAME], PreClosed + 1, 5000) + end, + Key = [?FUNCTION_NAME, mtp_policy:convert(client_ipv4, {127, 0, 0, 1})], + %% Open 2 connections, make sure 3rd one will be rejected + Cli10 = mtp_test_client:connect(Host, Port, Secret, DcId, mtp_secure), + Cli11 = ping(Cli10), + ?assertEqual(1, mtp_policy_counter:get(Key)), + Cli20 = mtp_test_client:connect(Host, Port, Secret, DcId, mtp_secure), + _Cli21 = ping(Cli20), + ?assertEqual(2, mtp_policy_counter:get(Key)), + Cli31 = mtp_test_client:connect(Host, Port, Secret, DcId, mtp_secure), + ?assertError({badmatch, {error, closed}}, ping(Cli31)), + ?assertEqual( + 1, mtp_test_metric:get_tags( + count, [?APP, protocol_error, total], [?FUNCTION_NAME, policy_error])), + ?assertEqual(2, mtp_policy_counter:get(Key)), + %% Close 1st connection and try to connect again. This should work. + SureClose(Cli11), + ?assertEqual(1, mtp_policy_counter:get(Key)), + Cli40 = mtp_test_client:connect(Host, Port, Secret, DcId, mtp_secure), + _Cli41 = ping(Cli40), + ?assertEqual(2, mtp_policy_counter:get(Key)), + ok. + +%% @doc tests that connections to whitelistsed domains are allowed and not from the list disallowed +policy_whitelist_case({pre, Cfg}) -> + Cfg1 = setup_single(?FUNCTION_NAME, 10000 + ?LINE, #{}, Cfg), + %% Allow max 2 connections from IP + Domain = <<"allowed.example.com">>, + ok = mtp_policy_table:add(domain_whitelist, tls_domain, Domain), + set_env([{policy, [{in_table, tls_domain, domain_whitelist}]}], + [{domain, Domain} | Cfg1]); +policy_whitelist_case({post, Cfg}) -> + stop_single(Cfg), + reset_env(Cfg); +policy_whitelist_case(Cfg) when is_list(Cfg) -> + DcId = ?config(dc_id, Cfg), + Host = ?config(mtp_host, Cfg), + Port = ?config(mtp_port, Cfg), + Secret = ?config(mtp_secret, Cfg), + Domain = ?config(domain, Cfg), + Cli01 = mtp_test_client:connect(Host, Port, Secret, DcId, + {mtp_fake_tls, Domain}), + _Cli02 = ping(Cli01), + ?assertError({badmatch, {error, closed}}, + begin + Cli11 = mtp_test_client:connect(Host, Port, Secret, DcId, + {mtp_fake_tls, <<"not-", Domain/binary>>}), + ping(Cli11) + end), + ?assertEqual( + 1, mtp_test_metric:get_tags( + count, [?APP, protocol_error, total], [?FUNCTION_NAME, policy_error])), + ok. + %% Helpers setup_single(Name, MtpPort, DcCfg0, Cfg) -> @@ -486,3 +561,10 @@ reset_env(Cfg) -> {ok, Val} -> application:set_env(mtproto_proxy, K, Val) end || {K, V} <- OldEnv]. + +ping(Cli0) -> + Data = crypto:strong_rand_bytes(64), + Cli1 = mtp_test_client:send(Data, Cli0), + {ok, Packet, Cli2} = mtp_test_client:recv_packet(Cli1, 1000), + ?assertEqual(Data, Packet), + Cli2.