librz/util/pj: Fix JSON depth limit handling (#6533)

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 */

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 };

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;
 }