fix: websocket's max_connection not work

This commit is contained in:
Zhongwen Deng 2022-04-27 13:58:38 +08:00
parent 09b6648567
commit ab9e93e8a7
2 changed files with 131 additions and 97 deletions

View File

@ -1515,7 +1515,7 @@ base_listener() ->
)}, )},
{"acceptors", {"acceptors",
sc( sc(
integer(), pos_integer(),
#{ #{
default => 16, default => 16,
desc => ?DESC(base_listener_acceptors) desc => ?DESC(base_listener_acceptors)
@ -1523,7 +1523,7 @@ base_listener() ->
)}, )},
{"max_connections", {"max_connections",
sc( sc(
hoconsc:union([infinity, integer()]), hoconsc:union([infinity, pos_integer()]),
#{ #{
default => infinity, default => infinity,
desc => ?DESC(base_listener_max_connections) desc => ?DESC(base_listener_max_connections)

View File

@ -272,51 +272,18 @@ check_origin_header(Req, #{listener := {Type, Listener}} = Opts) ->
false -> ok false -> ok
end. end.
websocket_init([ websocket_init([Req, Opts]) ->
Req, #{zone := Zone, limiter := LimiterCfg, listener := {Type, Listener}} = Opts,
#{zone := Zone, limiter := LimiterCfg, listener := {Type, Listener}} = Opts case check_max_connection(Type, Listener) of
]) -> allow ->
{Peername, Peercert} = {Peername, PeerCert} = get_peer_info(Type, Listener, Req, Opts),
case
emqx_config:get_listener_conf(Type, Listener, [proxy_protocol]) andalso
maps:get(proxy_header, Req)
of
#{src_address := SrcAddr, src_port := SrcPort, ssl := SSL} ->
SourceName = {SrcAddr, SrcPort},
%% Notice: Only CN is available in Proxy Protocol V2 additional info
SourceSSL =
case maps:get(cn, SSL, undefined) of
undeined -> nossl;
CN -> [{pp2_ssl_cn, CN}]
end,
{SourceName, SourceSSL};
#{src_address := SrcAddr, src_port := SrcPort} ->
SourceName = {SrcAddr, SrcPort},
{SourceName, nossl};
_ ->
{get_peer(Req, Opts), cowboy_req:cert(Req)}
end,
Sockname = cowboy_req:sock(Req), Sockname = cowboy_req:sock(Req),
WsCookie = WsCookie = get_ws_cookie(Req),
try
cowboy_req:parse_cookies(Req)
catch
error:badarg ->
?SLOG(error, #{msg => "bad_cookie"}),
undefined;
Error:Reason ->
?SLOG(error, #{
msg => "failed_to_parse_cookie",
exception => Error,
reason => Reason
}),
undefined
end,
ConnInfo = #{ ConnInfo = #{
socktype => ws, socktype => ws,
peername => Peername, peername => Peername,
sockname => Sockname, sockname => Sockname,
peercert => Peercert, peercert => PeerCert,
ws_cookie => WsCookie, ws_cookie => WsCookie,
conn_mod => ?MODULE conn_mod => ?MODULE
}, },
@ -331,28 +298,12 @@ websocket_init([
ParseState = emqx_frame:initial_parse_state(FrameOpts), ParseState = emqx_frame:initial_parse_state(FrameOpts),
Serialize = emqx_frame:serialize_opts(), Serialize = emqx_frame:serialize_opts(),
Channel = emqx_channel:init(ConnInfo, Opts), Channel = emqx_channel:init(ConnInfo, Opts),
GcState = GcState = get_force_gc(Zone),
case emqx_config:get_zone_conf(Zone, [force_gc]) of StatsTimer = get_stats_enable(Zone),
#{enable := false} -> undefined;
GcPolicy -> emqx_gc:init(GcPolicy)
end,
StatsTimer =
case emqx_config:get_zone_conf(Zone, [stats, enable]) of
true -> undefined;
false -> disabled
end,
%% MQTT Idle Timeout %% MQTT Idle Timeout
IdleTimeout = emqx_channel:get_mqtt_conf(Zone, idle_timeout), IdleTimeout = emqx_channel:get_mqtt_conf(Zone, idle_timeout),
IdleTimer = start_timer(IdleTimeout, idle_timeout), IdleTimer = start_timer(IdleTimeout, idle_timeout),
case tune_heap_size(Channel),
emqx_config:get_zone_conf(
emqx_channel:info(zone, Channel),
[force_shutdown]
)
of
#{enable := false} -> ok;
ShutdownPolicy -> emqx_misc:tune_heap_size(ShutdownPolicy)
end,
emqx_logger:set_metadata_peername(esockd:format(Peername)), emqx_logger:set_metadata_peername(esockd:format(Peername)),
{ok, {ok,
#state{ #state{
@ -374,7 +325,70 @@ websocket_init([
limiter_timer = undefined, limiter_timer = undefined,
limiter_cache = queue:new() limiter_cache = queue:new()
}, },
hibernate}. hibernate};
{denny, Reason} ->
{stop, Reason}
end.
tune_heap_size(Channel) ->
case
emqx_config:get_zone_conf(
emqx_channel:info(zone, Channel),
[force_shutdown]
)
of
#{enable := false} -> ok;
ShutdownPolicy -> emqx_misc:tune_heap_size(ShutdownPolicy)
end.
get_stats_enable(Zone) ->
case emqx_config:get_zone_conf(Zone, [stats, enable]) of
true -> undefined;
false -> disabled
end.
get_force_gc(Zone) ->
case emqx_config:get_zone_conf(Zone, [force_gc]) of
#{enable := false} -> undefined;
GcPolicy -> emqx_gc:init(GcPolicy)
end.
get_ws_cookie(Req) ->
try
cowboy_req:parse_cookies(Req)
catch
error:badarg ->
?SLOG(error, #{msg => "bad_cookie"}),
undefined;
Error:Reason ->
?SLOG(error, #{
msg => "failed_to_parse_cookie",
exception => Error,
reason => Reason
}),
undefined
end.
get_peer_info(Type, Listener, Req, Opts) ->
case
emqx_config:get_listener_conf(Type, Listener, [proxy_protocol]) andalso
maps:get(proxy_header, Req)
of
#{src_address := SrcAddr, src_port := SrcPort, ssl := SSL} ->
SourceName = {SrcAddr, SrcPort},
%% Notice: Only CN is available in Proxy Protocol V2 additional info
SourceSSL =
case maps:get(cn, SSL, undefined) of
undeined -> nossl;
CN -> [{pp2_ssl_cn, CN}]
end,
{SourceName, SourceSSL};
#{src_address := SrcAddr, src_port := SrcPort} ->
SourceName = {SrcAddr, SrcPort},
{SourceName, nossl};
_ ->
{get_peer(Req, Opts), cowboy_req:cert(Req)}
end.
websocket_handle({binary, Data}, State) when is_list(Data) -> websocket_handle({binary, Data}, State) when is_list(Data) ->
websocket_handle({binary, iolist_to_binary(Data)}, State); websocket_handle({binary, iolist_to_binary(Data)}, State);
@ -1000,6 +1014,26 @@ get_peer(Req, #{listener := {Type, Listener}}) ->
_:_ -> {Addr, PeerPort} _:_ -> {Addr, PeerPort}
end. end.
check_max_connection(Type, Listener) ->
case emqx_config:get_listener_conf(Type, Listener, [max_connections]) of
infinity ->
allow;
Max ->
MatchSpec = [{{'_', emqx_ws_connection}, [], [true]}],
Curr = ets:select_count(emqx_channel_conn, MatchSpec),
case Curr >= Max of
false ->
allow;
true ->
Reason = #{
max => Max,
current => Curr,
msg => "websocket_max_connections_limited"
},
?SLOG(warning, Reason),
{denny, Reason}
end
end.
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------
%% For CT tests %% For CT tests
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------