fix: bound py_state and harden stream/log source builders

py_state gains an optional max_state_entries cap (default infinity, so
existing behavior is unchanged) enforced with atomic admission
(insert_new + update_counter, ets:take on remove) so Python-driven
state_set can't exhaust node memory. A reserved sentinel row holds the
count and is protected from caller corruption and hidden from keys/fetch.

The py:stream and setup_logging helpers that build Python source now
validate module/function/kwarg names as identifiers (rejecting injection
at positions where quoting is meaningless) and escape string-literal
values, including newlines and other control bytes.

Adds py_state_SUITE (cap accounting + sentinel) and
py_stream_SUITE:test_stream_rejects_injection.

Author: benoitc

CHANGELOG.md

@@ -4,6 +4,13 @@
 
 ### Security
 
+- **Bounded shared state + safe stream/log builders** - `py_state` gained an optional
+  `max_state_entries` cap (default `infinity`, unchanged behavior) enforced with atomic
+  admission so Python-driven `state_set` can't exhaust node memory, and its size counter
+  is protected from corruption. The `py:stream` and logging helpers that build Python
+  source now strictly validate module/function/kwarg names as identifiers (rejecting
+  injection at positions where quoting is meaningless) and escape string-literal values
+  including control characters.
 - **Validated event-loop fd handles** - The asyncio reader/writer integration no longer
   hands Python a raw `fd_resource` pointer as an integer key. Each handle is an opaque id
   validated against a registry on every use, so a stale, duplicate, or fabricated id is a

src/py.erl

@@ -413,8 +413,8 @@ stream(Module, Func, Args, Kwargs) when map_size(Kwargs) == 0 ->
 stream(Module, Func, Args, Kwargs) ->
     %% With kwargs - use eval approach
     Ctx = py_context_router:get_context(),
-    ModuleBin = ensure_binary(Module),
-    FuncBin = ensure_binary(Func),
+    ModuleBin = valid_py_module(ensure_binary(Module)),
+    FuncBin = valid_py_ident(ensure_binary(Func)),
     KwargsCode = format_kwargs(Kwargs),
     ArgsCode = format_args(Args),
     Code = iolist_to_binary([
@@ -445,16 +445,16 @@ format_args(Args) ->
 %% @private Format a single argument
 format_arg(A) when is_integer(A) -> integer_to_binary(A);
 format_arg(A) when is_float(A) -> float_to_binary(A);
-format_arg(A) when is_binary(A) -> <<"'", A/binary, "'">>;
-format_arg(A) when is_atom(A) -> <<"'", (atom_to_binary(A))/binary, "'">>;
+format_arg(A) when is_binary(A) -> <<"'", (escape_py_literal(A))/binary, "'">>;
+format_arg(A) when is_atom(A) -> <<"'", (escape_py_literal(atom_to_binary(A)))/binary, "'">>;
 format_arg(A) when is_list(A) -> iolist_to_binary([<<"[">>, format_args(A), <<"]">>]);
 format_arg(_) -> <<"None">>.
 
 %% @private Format kwargs for Python code
 format_kwargs(Kwargs) when map_size(Kwargs) == 0 -> <<>>;
 format_kwargs(Kwargs) ->
     KwList = maps:fold(fun(K, V, Acc) ->
-        KB = if is_atom(K) -> atom_to_binary(K); is_binary(K) -> K end,
+        KB = valid_py_ident(if is_atom(K) -> atom_to_binary(K); is_binary(K) -> K end),
         [<<KB/binary, "=", (format_arg(V))/binary>> | Acc]
     end, [], Kwargs),
     iolist_to_binary([<<", ">>, lists:join(<<", ">>, KwList)]).
@@ -543,7 +543,9 @@ stream_start(Module, Func, Args, Opts) ->
     {ok, Ref}.
 
 %% @private Run the streaming via Python code
-stream_run_python(ModuleBin, FuncBin, RefHash) ->
+stream_run_python(ModuleBin0, FuncBin0, RefHash) ->
+    ModuleBin = valid_py_module(ModuleBin0),
+    FuncBin = valid_py_ident(FuncBin0),
     RefHashBin = integer_to_binary(RefHash),
     %% Build Python code that streams values using callbacks
     Code = iolist_to_binary([
@@ -1070,6 +1072,51 @@ escape_python_string(Str) ->
                      (C) -> [C]
                   end, Str).
 
+%% @private Escape a binary for safe embedding inside a single-quoted Python
+%% string literal: quote, backslash, and newline/CR/tab/other control bytes that
+%% would otherwise break out of or corrupt the literal.
+escape_py_literal(Bin) when is_binary(Bin) ->
+    << <<(escape_py_byte(B))/binary>> || <<B>> <= Bin >>.
+
+escape_py_byte($') -> <<"\\'">>;
+escape_py_byte($\\) -> <<"\\\\">>;
+escape_py_byte($\n) -> <<"\\n">>;
+escape_py_byte($\r) -> <<"\\r">>;
+escape_py_byte($\t) -> <<"\\t">>;
+escape_py_byte(B) when B < 16#20; B =:= 16#7f ->
+    list_to_binary(io_lib:format("\\x~2.16.0b", [B]));
+escape_py_byte(B) -> <<B>>.
+
+%% @private Validate a Python identifier ([A-Za-z_][A-Za-z0-9_]*). Crashes on a
+%% non-conforming value so an attacker-controlled module/func/kwarg name can't
+%% inject code at an identifier position (where quoting is meaningless).
+valid_py_ident(Bin) when is_binary(Bin), byte_size(Bin) > 0 ->
+    case ident_ok(Bin, first) of
+        true -> Bin;
+        false -> error({invalid_python_identifier, Bin})
+    end;
+valid_py_ident(Other) ->
+    error({invalid_python_identifier, Other}).
+
+%% @private Validate a dotted Python module path (each segment an identifier).
+valid_py_module(Bin) when is_binary(Bin), byte_size(Bin) > 0 ->
+    Segments = binary:split(Bin, <<".">>, [global]),
+    lists:foreach(fun valid_py_ident/1, Segments),
+    Bin;
+valid_py_module(Other) ->
+    error({invalid_python_identifier, Other}).
+
+ident_ok(<<>>, first) -> false;   %% empty segment (leading/trailing/double dot)
+ident_ok(<<>>, rest) -> true;
+ident_ok(<<C, Rest/binary>>, first)
+  when (C >= $A andalso C =< $Z); (C >= $a andalso C =< $z); C =:= $_ ->
+    ident_ok(Rest, rest);
+ident_ok(<<C, Rest/binary>>, rest)
+  when (C >= $A andalso C =< $Z); (C >= $a andalso C =< $z);
+       (C >= $0 andalso C =< $9); C =:= $_ ->
+    ident_ok(Rest, rest);
+ident_ok(_, _) -> false.
+
 %% @doc Deactivate the current virtual environment.
 %% Restores sys.path to its original state.
 -spec deactivate_venv() -> ok | {error, term()}.
@@ -1262,13 +1309,13 @@ configure_logging(Opts) ->
             iolist_to_binary([
                 "__import__('erlang').setup_logging(",
                 integer_to_binary(LevelInt),
-                ", '", F, "')"
+                ", '", escape_py_literal(F), "')"
             ]);
         F when is_list(F) ->
             iolist_to_binary([
                 "__import__('erlang').setup_logging(",
                 integer_to_binary(LevelInt),
-                ", '", F, "')"
+                ", '", escape_py_literal(iolist_to_binary(F)), "')"
             ])
     end,
     case eval(Code) of

src/py_state.erl

@@ -72,6 +72,11 @@
 
 -define(TABLE, py_state).
 
+%% Reserved sentinel row holding the live entry count, used by the optional size
+%% cap. User keys are rejected from this slot so callers (Erlang or Python) can't
+%% corrupt the accounting.
+-define(SIZE_KEY, '$py_state_size$').
+
 %%% ============================================================================
 %%% API
 %%% ============================================================================
@@ -112,28 +117,70 @@ register_callbacks() ->
 
 %% @doc Fetch a value from the shared state.
 -spec fetch(Key :: term()) -> {ok, term()} | {error, not_found}.
+fetch(?SIZE_KEY) -> {error, not_found};
 fetch(Key) ->
     case ets:lookup(?TABLE, Key) of
         [{_, Value}] -> {ok, Value};
         [] -> {error, not_found}
     end.
 
 %% @doc Store a value in the shared state.
--spec store(Key :: term(), Value :: term()) -> ok.
+-spec store(Key :: term(), Value :: term()) -> ok | {error, full | reserved_key}.
+store(?SIZE_KEY, _Value) ->
+    {error, reserved_key};
 store(Key, Value) ->
-    ets:insert(?TABLE, {Key, Value}),
-    ok.
+    case max_entries() of
+        infinity ->
+            ets:insert(?TABLE, {Key, Value}),
+            ok;
+        Max ->
+            %% Atomic admission: only genuinely new keys consume capacity, and the
+            %% reserve/rollback uses ets:update_counter so there is no TOCTOU race
+            %% on the public, write-concurrent table. Overwrites don't change count.
+            case ets:insert_new(?TABLE, {Key, Value}) of
+                true ->
+                    Count = ets:update_counter(?TABLE, ?SIZE_KEY, {2, 1}, {?SIZE_KEY, 0}),
+                    case Count > Max of
+                        true ->
+                            ets:delete(?TABLE, Key),
+                            ets:update_counter(?TABLE, ?SIZE_KEY, {2, -1}, {?SIZE_KEY, 0}),
+                            {error, full};
+                        false ->
+                            ok
+                    end;
+                false ->
+                    ets:insert(?TABLE, {Key, Value}),
+                    ok
+            end
+    end.
 
 %% @doc Remove a key from the shared state.
 -spec remove(Key :: term()) -> ok.
+remove(?SIZE_KEY) ->
+    ok;
 remove(Key) ->
-    ets:delete(?TABLE, Key),
-    ok.
+    case max_entries() of
+        infinity ->
+            ets:delete(?TABLE, Key),
+            ok;
+        _Max ->
+            %% Decrement only when a real user key was actually present (ets:take so
+            %% a missing-key remove can't drift the counter negative).
+            case ets:take(?TABLE, Key) of
+                [_] ->
+                    ets:update_counter(?TABLE, ?SIZE_KEY, {2, -1}, {?SIZE_KEY, 0}),
+                    ok;
+                [] ->
+                    ok
+            end
+    end.
 
 %% @doc Get all keys in the shared state.
 -spec keys() -> [term()].
 keys() ->
-    ets:foldl(fun({K, _}, Acc) -> [K | Acc] end, [], ?TABLE).
+    ets:foldl(fun({?SIZE_KEY, _}, Acc) -> Acc;
+                 ({K, _}, Acc) -> [K | Acc]
+              end, [], ?TABLE).
 
 %% @doc Clear all entries from the shared state.
 -spec clear() -> ok.
@@ -148,6 +195,8 @@ incr(Key) ->
 
 %% @doc Atomically increment a counter by Amount. Initializes to Amount if not exists.
 -spec incr(Key :: term(), Amount :: integer()) -> integer().
+incr(?SIZE_KEY, _Amount) ->
+    error(reserved_key);
 incr(Key, Amount) ->
     try
         ets:update_counter(?TABLE, Key, {2, Amount})
@@ -168,6 +217,12 @@ decr(Key) ->
 decr(Key, Amount) ->
     incr(Key, -Amount).
 
+%% @private Configured entry cap. `infinity' (the default) preserves the previous
+%% unbounded behavior; set application env `max_state_entries' to a positive
+%% integer to bound memory growth from Python-driven state_set calls.
+max_entries() ->
+    application:get_env(erlang_python, max_state_entries, infinity).
+
 %%% ============================================================================
 %%% Callback wrappers (for Python access)
 %%% ============================================================================
@@ -181,8 +236,10 @@ state_get_callback([Key]) ->
 
 %% @private
 state_set_callback([Key, Value]) ->
-    store(Key, Value),
-    none.
+    case store(Key, Value) of
+        ok -> none;
+        {error, Reason} -> {error, Reason}
+    end.
 
 %% @private
 state_delete_callback([Key]) ->

test/py_state_SUITE.erl

@@ -0,0 +1,91 @@
+%%% @doc Common Test suite for py_state shared-state store, focused on the
+%%% optional entry cap and its accounting (memory-exhaustion resistance).
+-module(py_state_SUITE).
+
+-include_lib("common_test/include/ct.hrl").
+
+-export([
+    all/0,
+    init_per_suite/1,
+    end_per_suite/1,
+    init_per_testcase/2,
+    end_per_testcase/2
+]).
+
+-export([
+    cap_disabled_test/1,
+    cap_enforced_test/1,
+    cap_accounting_test/1,
+    sentinel_protected_test/1
+]).
+
+-define(SIZE_KEY, '$py_state_size$').
+
+all() ->
+    [cap_disabled_test, cap_enforced_test, cap_accounting_test, sentinel_protected_test].
+
+init_per_suite(Config) ->
+    {ok, _} = application:ensure_all_started(erlang_python),
+    Config.
+
+end_per_suite(_Config) ->
+    application:stop(erlang_python),
+    ok.
+
+init_per_testcase(_TC, Config) ->
+    application:unset_env(erlang_python, max_state_entries),
+    py_state:clear(),
+    Config.
+
+end_per_testcase(_TC, _Config) ->
+    application:unset_env(erlang_python, max_state_entries),
+    py_state:clear(),
+    ok.
+
+%% @doc Default (infinity) preserves the previous unbounded behavior.
+cap_disabled_test(_Config) ->
+    [ok = py_state:store(N, N) || N <- lists:seq(1, 1000)],
+    {ok, 500} = py_state:fetch(500),
+    ok.
+
+%% @doc A finite cap rejects new keys beyond the limit; overwrites don't consume
+%% capacity.
+cap_enforced_test(_Config) ->
+    application:set_env(erlang_python, max_state_entries, 5),
+    [ok = py_state:store({k, N}, N) || N <- lists:seq(1, 5)],
+    {error, full} = py_state:store({k, 6}, 6),
+    %% Overwriting an existing key does NOT consume capacity.
+    ok = py_state:store({k, 3}, 30),
+    {ok, 30} = py_state:fetch({k, 3}),
+    {error, full} = py_state:store({k, 7}, 7),
+    ok.
+
+%% @doc Removing frees capacity; removing a missing key must not drift the count
+%% (which would let the cap be bypassed or never trip).
+cap_accounting_test(_Config) ->
+    application:set_env(erlang_python, max_state_entries, 3),
+    ok = py_state:store(a, 1),
+    ok = py_state:store(b, 2),
+    ok = py_state:store(c, 3),
+    {error, full} = py_state:store(d, 4),
+    %% Remove a missing key repeatedly: no phantom capacity freed.
+    [py_state:remove(missing) || _ <- lists:seq(1, 10)],
+    {error, full} = py_state:store(d, 4),
+    %% Remove a real key: frees exactly one slot.
+    ok = py_state:remove(a),
+    ok = py_state:store(d, 4),
+    {error, full} = py_state:store(e, 5),
+    ok.
+
+%% @doc The internal size sentinel can't be written or counter-corrupted by a
+%% caller, and is hidden from keys/0 and fetch/1.
+sentinel_protected_test(_Config) ->
+    application:set_env(erlang_python, max_state_entries, 2),
+    {error, reserved_key} = py_state:store(?SIZE_KEY, 9999),
+    {'EXIT', _} = (catch py_state:incr(?SIZE_KEY, 100)),
+    ok = py_state:store(x, 1),
+    ok = py_state:store(y, 2),
+    {error, full} = py_state:store(z, 3),
+    {error, not_found} = py_state:fetch(?SIZE_KEY),
+    false = lists:member(?SIZE_KEY, py_state:keys()),
+    ok.

test/py_stream_SUITE.erl

@@ -17,7 +17,8 @@
     test_stream_cancel/1,
     test_stream_error/1,
     test_stream_empty/1,
-    test_stream_large/1
+    test_stream_large/1,
+    test_stream_rejects_injection/1
 ]).
 
 all() ->
@@ -29,7 +30,8 @@ all() ->
         test_stream_cancel,
         test_stream_error,
         test_stream_empty,
-        test_stream_large
+        test_stream_large,
+        test_stream_rejects_injection
     ].
 
 init_per_suite(Config) ->
@@ -40,6 +42,18 @@ end_per_suite(_Config) ->
     ok = application:stop(erlang_python),
     ok.
 
+%% @doc A module/func name (or kwarg key) that isn't a valid Python identifier
+%% must be rejected, not interpolated into the generated source where it could
+%% inject code. Regression for the stream source-builder hardening.
+test_stream_rejects_injection(_Config) ->
+    {'EXIT', {{invalid_python_identifier, _}, _}} =
+        (catch py:stream(<<"os'); __import__('os').system('x">>, <<"walk">>, [], #{k => 1})),
+    {'EXIT', {{invalid_python_identifier, _}, _}} =
+        (catch py:stream(<<"math">>, <<"sqrt'); evil(">>, [], #{k => 1})),
+    {'EXIT', {{invalid_python_identifier, _}, _}} =
+        (catch py:stream(<<"math">>, <<"sqrt">>, [], #{<<"bad key)">> => 1})),
+    ok.
+
 %% Helper to collect all stream events
 collect_stream(Ref) ->
     collect_stream(Ref, [], 5000).