Merge pull request #7796 from zhongwencool/fix-ws-max-connection-not-work

fix: websocket's max_connection not work
This commit is contained in:
zhongwencool 2022-04-27 18:30:16 +08:00 committed by GitHub
commit 737abc5700
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
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,78 +272,65 @@ 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 Sockname = cowboy_req:sock(Req),
emqx_config:get_listener_conf(Type, Listener, [proxy_protocol]) andalso WsCookie = get_ws_cookie(Req),
maps:get(proxy_header, Req) ConnInfo = #{
of socktype => ws,
#{src_address := SrcAddr, src_port := SrcPort, ssl := SSL} -> peername => Peername,
SourceName = {SrcAddr, SrcPort}, sockname => Sockname,
%% Notice: Only CN is available in Proxy Protocol V2 additional info peercert => PeerCert,
SourceSSL = ws_cookie => WsCookie,
case maps:get(cn, SSL, undefined) of conn_mod => ?MODULE
undeined -> nossl; },
CN -> [{pp2_ssl_cn, CN}] Limiter = emqx_limiter_container:get_limiter_by_names(
end, [?LIMITER_BYTES_IN, ?LIMITER_MESSAGE_IN], LimiterCfg
{SourceName, SourceSSL}; ),
#{src_address := SrcAddr, src_port := SrcPort} -> MQTTPiggyback = get_ws_opts(Type, Listener, mqtt_piggyback),
SourceName = {SrcAddr, SrcPort}, FrameOpts = #{
{SourceName, nossl}; strict_mode => emqx_config:get_zone_conf(Zone, [mqtt, strict_mode]),
_ -> max_size => emqx_config:get_zone_conf(Zone, [mqtt, max_packet_size])
{get_peer(Req, Opts), cowboy_req:cert(Req)} },
end, ParseState = emqx_frame:initial_parse_state(FrameOpts),
Sockname = cowboy_req:sock(Req), Serialize = emqx_frame:serialize_opts(),
WsCookie = Channel = emqx_channel:init(ConnInfo, Opts),
try GcState = get_force_gc(Zone),
cowboy_req:parse_cookies(Req) StatsTimer = get_stats_enable(Zone),
catch %% MQTT Idle Timeout
error:badarg -> IdleTimeout = emqx_channel:get_mqtt_conf(Zone, idle_timeout),
?SLOG(error, #{msg => "bad_cookie"}), IdleTimer = start_timer(IdleTimeout, idle_timeout),
undefined; tune_heap_size(Channel),
Error:Reason -> emqx_logger:set_metadata_peername(esockd:format(Peername)),
?SLOG(error, #{ {ok,
msg => "failed_to_parse_cookie", #state{
exception => Error, peername = Peername,
reason => Reason sockname = Sockname,
}), sockstate = running,
undefined mqtt_piggyback = MQTTPiggyback,
end, limiter = Limiter,
ConnInfo = #{ parse_state = ParseState,
socktype => ws, serialize = Serialize,
peername => Peername, channel = Channel,
sockname => Sockname, gc_state = GcState,
peercert => Peercert, postponed = [],
ws_cookie => WsCookie, stats_timer = StatsTimer,
conn_mod => ?MODULE idle_timeout = IdleTimeout,
}, idle_timer = IdleTimer,
Limiter = emqx_limiter_container:get_limiter_by_names( zone = Zone,
[?LIMITER_BYTES_IN, ?LIMITER_MESSAGE_IN], LimiterCfg listener = {Type, Listener},
), limiter_timer = undefined,
MQTTPiggyback = get_ws_opts(Type, Listener, mqtt_piggyback), limiter_cache = queue:new()
FrameOpts = #{ },
strict_mode => emqx_config:get_zone_conf(Zone, [mqtt, strict_mode]), hibernate};
max_size => emqx_config:get_zone_conf(Zone, [mqtt, max_packet_size]) {denny, Reason} ->
}, {stop, Reason}
ParseState = emqx_frame:initial_parse_state(FrameOpts), end.
Serialize = emqx_frame:serialize_opts(),
Channel = emqx_channel:init(ConnInfo, Opts), tune_heap_size(Channel) ->
GcState =
case emqx_config:get_zone_conf(Zone, [force_gc]) of
#{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
IdleTimeout = emqx_channel:get_mqtt_conf(Zone, idle_timeout),
IdleTimer = start_timer(IdleTimeout, idle_timeout),
case case
emqx_config:get_zone_conf( emqx_config:get_zone_conf(
emqx_channel:info(zone, Channel), emqx_channel:info(zone, Channel),
@ -352,29 +339,56 @@ websocket_init([
of of
#{enable := false} -> ok; #{enable := false} -> ok;
ShutdownPolicy -> emqx_misc:tune_heap_size(ShutdownPolicy) ShutdownPolicy -> emqx_misc:tune_heap_size(ShutdownPolicy)
end, end.
emqx_logger:set_metadata_peername(esockd:format(Peername)),
{ok, get_stats_enable(Zone) ->
#state{ case emqx_config:get_zone_conf(Zone, [stats, enable]) of
peername = Peername, true -> undefined;
sockname = Sockname, false -> disabled
sockstate = running, end.
mqtt_piggyback = MQTTPiggyback,
limiter = Limiter, get_force_gc(Zone) ->
parse_state = ParseState, case emqx_config:get_zone_conf(Zone, [force_gc]) of
serialize = Serialize, #{enable := false} -> undefined;
channel = Channel, GcPolicy -> emqx_gc:init(GcPolicy)
gc_state = GcState, end.
postponed = [],
stats_timer = StatsTimer, get_ws_cookie(Req) ->
idle_timeout = IdleTimeout, try
idle_timer = IdleTimer, cowboy_req:parse_cookies(Req)
zone = Zone, catch
listener = {Type, Listener}, error:badarg ->
limiter_timer = undefined, ?SLOG(error, #{msg => "bad_cookie"}),
limiter_cache = queue:new() undefined;
}, Error:Reason ->
hibernate}. ?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
%%-------------------------------------------------------------------- %%--------------------------------------------------------------------