From 31b9eb2f181ff9275ccbba787a8c1f991a70f69d Mon Sep 17 00:00:00 2001
From: Junxian Huang <huangjunxian6@hisilicon.com>
Date: Mon, 5 Aug 2024 19:59:19 +0800
Subject: [PATCH 05/27] hikptool/roce: Add check for total_block_num from FW

total_block_num is used to allocate memory for register array.
Check whether it is 0 to prevent calloc() error.

Fixes: 8b3b68347165 ("hikptool/roce: Add a common frame for hikptool roce register query")
Signed-off-by: Junxian Huang <huangjunxian6@hisilicon.com>
---
 net/roce/roce_ext_common/hikp_roce_ext_common.c | 6 ++++++
 net/roce/roce_scc/hikp_roce_scc.c               | 6 ++++++
 net/roce/roce_trp/hikp_roce_trp.c               | 6 ++++++
 3 files changed, 18 insertions(+)

diff --git a/net/roce/roce_ext_common/hikp_roce_ext_common.c b/net/roce/roce_ext_common/hikp_roce_ext_common.c
index 049b64a..bc3b883 100644
--- a/net/roce/roce_ext_common/hikp_roce_ext_common.c
+++ b/net/roce/roce_ext_common/hikp_roce_ext_common.c
@@ -122,6 +122,12 @@ static int hikp_roce_ext_get_res(enum roce_cmd_type cmd_type,
 
 	if (block_id == 0) {
 		res_head->total_block_num = roce_ext_res->head.total_block_num;
+		if (!res_head->total_block_num) {
+			printf("hikptool roce_%s total_block_num error!\n",
+				cmd_name);
+			ret = -EINVAL;
+			goto get_data_error;
+		}
 		reg->offset = (uint32_t *)calloc(res_head->total_block_num, sizeof(uint32_t));
 		reg->data = (uint32_t *)calloc(res_head->total_block_num, sizeof(uint32_t));
 		if ((reg->offset == NULL) || (reg->data == NULL)) {
diff --git a/net/roce/roce_scc/hikp_roce_scc.c b/net/roce/roce_scc/hikp_roce_scc.c
index 76c0ca6..0ecb8cf 100644
--- a/net/roce/roce_scc/hikp_roce_scc.c
+++ b/net/roce/roce_scc/hikp_roce_scc.c
@@ -146,6 +146,12 @@ static int hikp_roce_scc_get_total_data_num(struct roce_scc_head *res_head,
 	}
 
 	roce_scc_res = (struct roce_scc_res_param *)cmd_ret->rsp_data;
+	if (!roce_scc_res->head.total_block_num) {
+		printf("hikptool roce_scc total_block_num error!\n");
+		ret = -EINVAL;
+		goto get_data_error;
+	}
+
 	max_size = roce_scc_res->head.total_block_num * sizeof(uint32_t);
 	*offset = (uint32_t *)calloc(1, max_size);
 	*data = (uint32_t *)calloc(1, max_size);
diff --git a/net/roce/roce_trp/hikp_roce_trp.c b/net/roce/roce_trp/hikp_roce_trp.c
index 61f0511..486bbe3 100644
--- a/net/roce/roce_trp/hikp_roce_trp.c
+++ b/net/roce/roce_trp/hikp_roce_trp.c
@@ -166,6 +166,12 @@ static int hikp_roce_trp_get_total_data_num(struct roce_trp_head *res_head,
 	}
 
 	roce_trp_res = (struct roce_trp_res_param *)cmd_ret->rsp_data;
+	if (!roce_trp_res->head.total_block_num) {
+		printf("hikptool roce_trp total_block_num error!\n");
+		ret = -EINVAL;
+		goto get_data_error;
+	}
+
 	max_size = roce_trp_res->head.total_block_num * sizeof(uint32_t);
 	*offset = (uint32_t *)calloc(1, max_size);
 	*data = (uint32_t *)calloc(1, max_size);
-- 
2.45.0.windows.1