diff --git a/librz/include/rz_util/rz_pj.h b/librz/include/rz_util/rz_pj.h index 2e3f9eceb1..1601df671e 100644 --- a/librz/include/rz_util/rz_pj.h +++ b/librz/include/rz_util/rz_pj.h @@ -14,7 +14,7 @@ typedef struct pj_t { bool is_first; bool is_key; char braces[RZ_PRINT_JSON_DEPTH_LIMIT]; - int level; + st64 level; ///< Current nesting depth. 0 = top-level, >0 = inside braces, <0 = error (depth limit exceeded). } PJ; /* lifecycle */ diff --git a/librz/util/pj.c b/librz/util/pj.c index 7b61eb2ed1..5dcaa09dea 100644 --- a/librz/util/pj.c +++ b/librz/util/pj.c @@ -6,6 +6,9 @@ RZ_API void pj_raw(PJ *j, const char *msg) { rz_return_if_fail(j && msg); + if (j->level < 0) { + return; + } if (*msg) { rz_strbuf_append(&j->sb, msg); } @@ -13,6 +16,9 @@ RZ_API void pj_raw(PJ *j, const char *msg) { static void pj_comma(PJ *j) { rz_return_if_fail(j); + if (j->level < 0) { + return; + } if (!j->is_key) { if (!j->is_first) { pj_raw(j, ","); @@ -48,19 +54,28 @@ RZ_API void pj_reset(PJ *j) { } RZ_API char *pj_drain(PJ *pj) { - rz_return_val_if_fail(pj && pj->level == 0, NULL); + if (!pj) { + return NULL; + } + if (pj->level != 0) { + // error case, invalid json, return NULL + pj_free(pj); + return NULL; + } char *res = rz_strbuf_drain_nofree(&pj->sb); free(pj); return res; } RZ_API const char *pj_string(PJ *j) { - return j ? rz_strbuf_get(&j->sb) : NULL; + return (j && j->level >= 0) ? rz_strbuf_get(&j->sb) : NULL; } static PJ *pj_begin(PJ *j, char type) { if (j) { - if (!j || j->level >= RZ_PRINT_JSON_DEPTH_LIMIT) { + if (j->level < 0 || j->level >= RZ_PRINT_JSON_DEPTH_LIMIT) { + j->level = -1; + RZ_LOG_ERROR("pj: JSON maximum depth of %d exceeded\n", RZ_PRINT_JSON_DEPTH_LIMIT); return NULL; } char msg[2] = { type, 0 }; diff --git a/test/unit/test_pj.c b/test/unit/test_pj.c index 7047fd39e1..bea451cf8f 100644 --- a/test/unit/test_pj.c +++ b/test/unit/test_pj.c @@ -22,8 +22,53 @@ bool test_pj_reset() { mu_end; } +bool test_depth_limit() { + PJ *j = pj_new(); + for (int i = 0; i < RZ_PRINT_JSON_DEPTH_LIMIT + 1; i++) { + pj_o(j); + } + mu_assert_null(pj_string(j), "pj_string should be null after exceeding depth limit"); + + // writing after overflow is suppressed (level<0 path) + pj_ks(j, "key", "val"); + mu_assert_null(pj_string(j), "pj_string still null after writing in overflow state"); + + // pj_end is a nop after overflow (level<1 path) + pj_end(j); + mu_assert_null(pj_string(j), "pj_string still null after pj_end in overflow state"); + pj_free(j); + + mu_end; +} + +bool test_pj_drain() { + // pj_drain returns NULL and frees when level != 0 (overflow) + PJ *j = pj_new(); + for (int i = 0; i < RZ_PRINT_JSON_DEPTH_LIMIT + 1; i++) { + pj_o(j); + } + char *res = pj_drain(j); // j is consumed + mu_assert_null(res, "pj_drain should return NULL on overflow"); + + mu_assert_null(pj_drain(NULL), "pj_drain(NULL) should return NULL"); + + // pj_drain suceed when level=0 + j = pj_new(); + pj_o(j); + pj_ki(j, "a", 1); + pj_end(j); + res = pj_drain(j); + mu_assert_notnull(res, "pj_drain should succeed on valid json"); + mu_assert_streq(res, "{\"a\":1}", "pj_drain result"); + free(res); + + mu_end; +} + int all_tests() { mu_run_test(test_pj_reset); + mu_run_test(test_depth_limit); + mu_run_test(test_pj_drain); return tests_passed != tests_run; }