diff --git a/NetUtils.py b/NetUtils.py index f792a27288e..3db8a8bd425 100644 --- a/NetUtils.py +++ b/NetUtils.py @@ -170,7 +170,39 @@ def _object_hook(o: typing.Any) -> typing.Any: return o -decode = JSONDecoder(object_hook=_object_hook).decode +_decode = JSONDecoder(object_hook=_object_hook).decode + + +try: + from _speedups import check_json_depth as _check_depth +except ImportError: + def _check_depth(s: str, limit: int = 16) -> None: + depth = 1 + in_quotes = False + escape = False + for c in s: + if c == "\\" and not escape: + escape = True + continue + if c == '"' and not escape: + in_quotes = not in_quotes + continue + if not in_quotes: + if c in "[{": + depth += 1 + if depth > limit: + raise ValueError("JSON document too complex") + elif c in "}]": + depth -= 1 + escape = False + + if depth != 1: # free check + raise ValueError("JSON document malformed") + + +def decode(s: str) -> typing.Any: + _check_depth(s) # raises ValueError + return _decode(s) class Endpoint: diff --git a/_speedups.pyx b/_speedups.pyx index 2ad1a2953a2..8112a95846a 100644 --- a/_speedups.pyx +++ b/_speedups.pyx @@ -376,3 +376,27 @@ cdef class PlayerLocationProxy: count = self._store.sender_index[self._player].count for entry in self._store.entries[start:start+count]: yield entry.location, (entry.item, entry.receiver, entry.flags) + + +cpdef void check_json_depth(s: str, limit: int = 16): + cdef Py_ssize_t depth = 1 + cdef bool in_quotes = False + cdef bool escape = False + for c in s: + if c == "\\" and not escape: + escape = True + continue + if c == '"' and not escape: + in_quotes = not in_quotes + continue + if not in_quotes: + if c in "[{": + depth += 1 + if depth > limit: + raise ValueError("JSON document too complex") + elif c in "}]": + depth -= 1 + escape = False + + if depth != 1: # free check + raise ValueError("JSON document malformed") diff --git a/test/netutils/test_decode.py b/test/netutils/test_decode.py new file mode 100644 index 00000000000..4ac42a66124 --- /dev/null +++ b/test/netutils/test_decode.py @@ -0,0 +1,68 @@ +import unittest +from typing import Any + +from NetUtils import decode, encode + + +class DecodeDepthLimitTest(unittest.TestCase): + LIMIT = 16 + + @staticmethod + def make_data(depth: int = LIMIT, cmd: str = "Cmd") -> list[dict[str, Any]]: + arg: Any = [1] + for _ in range(depth - 4): + arg = [arg] + res = {"cmd": cmd, "arg": arg} + # [{"arg": [[...[1]...]]}] + # ^1 ^depth + return [res] + + @classmethod + def make_message(cls, depth: int = LIMIT, cmd: str = "Cmd") -> str: + return encode(cls.make_data(depth, cmd=cmd)) + + def test_below_limit(self) -> None: + data = self.make_data(depth=self.LIMIT - 1) + message = encode(data) + self.assertEqual(data, decode(message)) + + def test_at_limit(self) -> None: + data = self.make_data(depth=self.LIMIT) + message = encode(data) + self.assertEqual(data, decode(message)) + + def test_above_limit(self) -> None: + with self.assertRaises(ValueError): + decode(self.make_message(depth=self.LIMIT + 1)) + + def test_incomplete(self) -> None: + with self.assertRaises(ValueError): + decode(self.make_message()[:-1]) + + def test_invalid(self) -> None: + with self.assertRaises(ValueError): + decode(self.make_message().replace(":", ",")) + + def test_braces_in_str(self) -> None: + # should not raise + decode(self.make_message(cmd="[")) + decode(self.make_message(cmd="{")) + decode(self.make_message(cmd="}")) + decode(self.make_message(cmd="]")) + + def test_quote_in_str(self) -> None: + # should not raise + decode(self.make_message(cmd='"')) + + def test_bs_quote_in_str(self) -> None: + # should not raise + decode(self.make_message(cmd=r'\"')) + + def test_quoted_braces_in_str(self) -> None: + # should not raise + decode(self.make_message(cmd='"{["')) + + def test_escape(self) -> None: + # should not raise + decode(r"""["\"\\\/\b\f\n\r\t{"]""") + self.assertEqual("new\nline", decode(r'"new\u000Aline"'))