Skip to content

Commit d6045a6

Browse files
prefiksprefiks
committed
Add ability to specify max number of allowed elements inside stanza
Also support changing limits on already initialized parser
1 parent 3e2a3a9 commit d6045a6

3 files changed

Lines changed: 117 additions & 15 deletions

File tree

‎c_src/fxml_stream.c‎

Lines changed: 79 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
#include <expat.h>
2222

2323
#define PARSING_NOT_RESUMABLE XML_FALSE
24+
#define LIMIT_MAX ((ErlNifUInt64)-1)
2425

2526
#define ASSERT(x) if (!(x)) return 0
2627
#define PARSER_ASSERT(X, E) do { if (!(X)) { state->error = (E); XML_StopParser(state->parser, PARSING_NOT_RESUMABLE); return; } } while(0)
@@ -56,8 +57,10 @@ typedef struct {
5657
ErlNifEnv *send_env;
5758
ErlNifPid *pid;
5859
size_t depth;
59-
size_t size;
60-
size_t max_size;
60+
ErlNifUInt64 size;
61+
ErlNifUInt64 elements_count;
62+
ErlNifUInt64 max_size;
63+
ErlNifUInt64 max_elements;
6164
XML_Parser parser;
6265
xmlel_stack_t *elements_stack;
6366
attrs_list_t *xmlns_attrs;
@@ -256,6 +259,7 @@ static ERL_NIF_TERM str2bin(ErlNifEnv *env, const char *s)
256259
static void send_event(state_t *state, ERL_NIF_TERM el)
257260
{
258261
state->size = 0;
262+
state->elements_count = 0;
259263
if (state->gen_server) {
260264
enif_send(state->env, state->pid, state->send_env,
261265
enif_make_tuple2(state->send_env,
@@ -270,6 +274,7 @@ static void send_event(state_t *state, ERL_NIF_TERM el)
270274
static void send_all_state_event(state_t *state, ERL_NIF_TERM el)
271275
{
272276
state->size = 0;
277+
state->elements_count = 0;
273278
if (state->gen_server) {
274279
enif_send(state->env, state->pid, state->send_env,
275280
enif_make_tuple2(state->send_env,
@@ -311,6 +316,9 @@ void erlXML_StartElementHandler(state_t *state,
311316
if (state->error)
312317
return;
313318

319+
state->elements_count += 1;
320+
PARSER_ASSERT(state->elements_count < state->max_elements, "XML stanza is too big");
321+
314322
state->depth++;
315323

316324
while (atts[i])
@@ -851,6 +859,7 @@ static ERL_NIF_TERM reset_nif(ErlNifEnv* env, int argc,
851859
enif_clear_env(state->send_env);
852860

853861
state->size = 0;
862+
state->elements_count = 0;
854863
state->depth = 0;
855864
state->error = NULL;
856865

@@ -888,6 +897,9 @@ static ERL_NIF_TERM parse_element_nif(ErlNifEnv* env, int argc,
888897
if (!state)
889898
return enif_make_badarg(env);
890899

900+
state->max_elements = LIMIT_MAX;
901+
state->max_size = LIMIT_MAX;
902+
891903
state->send_env = env;
892904
state->use_maps = use_maps;
893905

@@ -962,11 +974,14 @@ static ERL_NIF_TERM parse_nif(ErlNifEnv* env, int argc,
962974
state->size += bin.size;
963975
state->env = env;
964976

977+
if (state->error) {
978+
send_error(state, str2bin(state->send_env, state->error));
979+
return argv[0];
980+
}
981+
965982
if (state->size >= state->max_size) {
966-
size_t size = state->size;
967-
send_error(state, str2bin(state->send_env, "XML stanza is too big"));
968-
/* Don't let send_event() to set size to zero */
969-
state->size = size;
983+
state->error = "XML stanza is too big";
984+
send_error(state, str2bin(state->send_env, state->error));
970985
} else {
971986
int res = XML_Parse(state->parser, (char *)bin.data, bin.size, 0);
972987
if (!res)
@@ -1059,17 +1074,68 @@ static ERL_NIF_TERM new_nif(ErlNifEnv* env, int argc,
10591074
ERL_NIF_TERM result = enif_make_resource(env, state);
10601075
enif_release_resource(state);
10611076

1062-
ErlNifUInt64 max_size;
1063-
if (enif_get_uint64(env, argv[1], &max_size))
1077+
ErlNifUInt64 max_size, max_elements;
1078+
int arity;
1079+
ERL_NIF_TERM *tuple_elements;
1080+
if (enif_get_uint64(env, argv[1], &max_size)) {
10641081
state->max_size = (size_t) max_size;
1065-
else if (!enif_compare(argv[1], enif_make_atom(env, "infinity")))
1066-
state->max_size = (size_t) - 1;
1067-
else
1082+
state->max_elements = LIMIT_MAX;
1083+
} else if (!enif_compare(argv[1], enif_make_atom(env, "infinity"))) {
1084+
state->max_size = LIMIT_MAX;
1085+
state->max_elements = LIMIT_MAX;
1086+
} else if (enif_get_tuple(env, argv[1], &arity, &tuple_elements) && arity == 2) {
1087+
if (enif_get_uint64(env, tuple_elements[0], &max_size)) {
1088+
state->max_size = max_size;
1089+
} else if (!enif_compare(tuple_elements[0], enif_make_atom(env, "infinity"))) {
1090+
state->max_size = LIMIT_MAX;
1091+
} else {
1092+
return enif_make_badarg(env);
1093+
}
1094+
1095+
if (enif_get_uint64(env, tuple_elements[1], &max_elements)) {
1096+
state->max_elements = max_elements;
1097+
} else if (!enif_compare(tuple_elements[1], enif_make_atom(env, "infinity"))) {
1098+
state->max_elements = LIMIT_MAX;
1099+
} else {
1100+
return enif_make_badarg(env);
1101+
}
1102+
} else
10681103
return enif_make_badarg(env);
10691104

10701105
return result;
10711106
}
10721107

1108+
static ERL_NIF_TERM change_limits_nif(ErlNifEnv* env, int argc,
1109+
const ERL_NIF_TERM argv[])
1110+
{
1111+
state_t *state = NULL;
1112+
1113+
if (argc != 3)
1114+
return enif_make_badarg(env);
1115+
1116+
if (!enif_get_resource(env, argv[0], parser_state_t, (void *) &state))
1117+
return enif_make_badarg(env);
1118+
1119+
ErlNifUInt64 max_size, max_elements;
1120+
if (enif_get_uint64(env, argv[1], &max_size)) {
1121+
} else if (!enif_compare(argv[1], enif_make_atom(env, "infinity"))) {
1122+
max_size = LIMIT_MAX;
1123+
} else
1124+
return enif_make_badarg(env);
1125+
1126+
if (enif_get_uint64(env, argv[2], &max_elements)) {
1127+
} else if (!enif_compare(argv[2], enif_make_atom(env, "infinity"))) {
1128+
max_elements = LIMIT_MAX;
1129+
} else
1130+
return enif_make_badarg(env);
1131+
1132+
state->max_size = max_size;
1133+
state->max_elements = max_elements;
1134+
1135+
return enif_make_atom(env, "ok");
1136+
}
1137+
1138+
10731139
static ErlNifFunc nif_funcs[] =
10741140
{
10751141
{"new", 2, new_nif},
@@ -1079,7 +1145,8 @@ static ErlNifFunc nif_funcs[] =
10791145
{"parse_element", 2, parse_element_nif},
10801146
{"reset", 1, reset_nif},
10811147
{"close", 1, close_nif},
1082-
{"change_callback_pid", 2, change_callback_pid_nif}
1148+
{"change_callback_pid", 2, change_callback_pid_nif},
1149+
{"change_limits", 3, change_limits_nif}
10831150
};
10841151

10851152
ERL_NIF_INIT(fxml_stream, nif_funcs, load, NULL, upgrade, NULL)

‎src/fxml_stream.erl‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
-on_load(init/0).
3131

3232
-export([new/1, new/2, new/3, parse/2, close/1, reset/1,
33-
change_callback_pid/2, parse_element/1, parse_element/2]).
33+
change_callback_pid/2, change_limits/3, parse_element/1, parse_element/2]).
3434

3535
-export([load_nif/0, load_nif/1]).
3636

@@ -77,12 +77,12 @@ load_nif(SOPath) ->
7777
new(CallbackPid) ->
7878
new(CallbackPid, infinity).
7979

80-
-spec new(pid(), non_neg_integer() | infinity) -> xml_stream_state().
80+
-spec new(pid(), non_neg_integer() | infinity | {non_neg_integer() | infinity, non_neg_integer() | infinity}) -> xml_stream_state().
8181

8282
new(_CallbackPid, _MaxSize) ->
8383
erlang:nif_error(nif_not_loaded).
8484

85-
-spec new(pid(), non_neg_integer() | infinity, list()) -> xml_stream_state().
85+
-spec new(pid(), non_neg_integer() | infinity | {non_neg_integer() | infinity, non_neg_integer() | infinity}, list()) -> xml_stream_state().
8686

8787
new(_CallbackPid, _MaxSize, _Options) ->
8888
erlang:nif_error(nif_not_loaded).
@@ -97,6 +97,10 @@ reset(_State) ->
9797
change_callback_pid(_State, _CallbackPid) ->
9898
erlang:nif_error(nif_not_loaded).
9999

100+
-spec change_limits(xml_stream_state(), non_neg_integer() | infinity, non_neg_integer() | infinity) -> ok.
101+
change_limits(_State, _NewMaxSize, _NewMaxElements) ->
102+
erlang:nif_error(nif_not_loaded).
103+
100104
-spec parse(xml_stream_state(), binary()) -> xml_stream_state().
101105

102106
parse(_State, _Data) ->

‎test/fxml_test.erl‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -419,6 +419,37 @@ too_big_with_data_after_test() ->
419419
<<"</foo></start>">>,
420420
[{xmlstreamerror, <<"XML stanza is too big">>}]
421421
}]).
422+
423+
too_much_elements_test() ->
424+
CallbackPid = spawn_link(fun() -> receiver([]) end),
425+
Stream0 = new(CallbackPid, {infinity, 3}),
426+
Stream1 = fxml_stream:parse(Stream0, <<"<a>">>),
427+
Stream2 = fxml_stream:parse(Stream1, <<"<b/>">>),
428+
Stream3 = fxml_stream:parse(Stream2, <<"<b><c/></b>">>),
429+
Stream4 = fxml_stream:parse(Stream3, <<"<b><c/><d/></b>">>),
430+
Stream5 = fxml_stream:parse(Stream4, <<"<b/>">>),
431+
?assertEqual([{xmlstreamstart, <<"a">>, []},
432+
{xmlstreamelement, #xmlel{name = <<"b">>}},
433+
{xmlstreamelement, #xmlel{name = <<"b">>, children = [#xmlel{name = <<"c">>}]}},
434+
{xmlstreamerror, <<"XML stanza is too big">>},
435+
{xmlstreamerror, <<"XML stanza is too big">>}],
436+
collect_events(CallbackPid)),
437+
close(Stream5).
438+
439+
runtime_limits_change_test() ->
440+
CallbackPid = spawn_link(fun() -> receiver([]) end),
441+
Stream0 = new(CallbackPid, {infinity, 3}),
442+
Stream1 = fxml_stream:parse(Stream0, <<"<a>">>),
443+
Stream2 = fxml_stream:parse(Stream1, <<"<b/>">>),
444+
Stream3 = fxml_stream:parse(Stream2, <<"<b><c/></b>">>),
445+
ok = fxml_stream:change_limits(Stream3, infinity, infinity),
446+
Stream4 = fxml_stream:parse(Stream3, <<"<b><c/><d/></b>">>),
447+
?assertEqual([{xmlstreamstart, <<"a">>, []},
448+
{xmlstreamelement, #xmlel{name = <<"b">>}},
449+
{xmlstreamelement, #xmlel{name = <<"b">>, children = [#xmlel{name = <<"c">>}]}},
450+
{xmlstreamelement, #xmlel{name = <<"b">>, children = [#xmlel{name = <<"c">>}, #xmlel{name = <<"d">>}]}}],
451+
collect_events(CallbackPid)),
452+
close(Stream4).
422453

423454
close_close_test() ->
424455
Stream = new(),

0 commit comments

Comments
 (0)