From e505a6320e9c0faeab7723a4ddc6d4db62494ab4 Mon Sep 17 00:00:00 2001
From: Wangjunqi123 <wangjunqi@kylinos.cn>
Date: Mon, 18 Nov 2024 10:02:58 +0800
Subject: [PATCH] agent:when sending a heartbeat request, judgement
 len(resp.body) before unmarshaling resp.body

---
 cmd/agent/service/heartbeat.go                | 43 +++++++++++++------
 cmd/server/webserver/handle/basicHandler.go   | 13 +++++-
 .../webserver/handle/customTopoHandler.go     |  2 +-
 3 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/cmd/agent/service/heartbeat.go b/cmd/agent/service/heartbeat.go
index 2e2dcde..0fda0a9 100644
--- a/cmd/agent/service/heartbeat.go
+++ b/cmd/agent/service/heartbeat.go
@@ -15,14 +15,21 @@ import (
 func SendHeartbeat() {
 	agent_addr := conf.Config().Topo.Agent_addr
 
+	global.ERManager.Wg.Add(1)
 	go func() {
+		global.ERManager.Wg.Done()
 		for {
-			err := sendHeartbeat(agent_addr)
-			if err != nil {
-				err = errors.Wrap(err, " ")
-				global.ERManager.ErrorTransmit("service", "error", err, false, false)
+			select {
+			case <-global.ERManager.GoCancelCtx.Done():
+				return
+			default:
+				err := sendHeartbeat(agent_addr)
+				if err != nil {
+					err = errors.Wrap(err, " ")
+					global.ERManager.ErrorTransmit("service", "error", err, false, false)
+				}
+				time.Sleep(time.Duration(conf.Config().Topo.Heartbeat) * time.Second)
 			}
-			time.Sleep(time.Duration(conf.Config().Topo.Heartbeat) * time.Second)
 		}
 	}()
 }
@@ -30,8 +37,7 @@ func SendHeartbeat() {
 func sendHeartbeat(addr string) error {
 	m_u_bytes, err := global.FileReadBytes(global.Agentuuid_filepath)
 	if err != nil {
-		err = errors.Wrap(err, " ")
-		return err
+		return errors.Wrap(err, " ")
 	}
 	type machineuuid struct {
 		Agentuuid string `json:"agent_uuid"`
@@ -54,26 +60,37 @@ func sendHeartbeat(addr string) error {
 	params := &httputils.Params{
 		Body: body,
 	}
+
 	resp, err := httputils.Post(url, params)
 	if err != nil {
-		err = errors.Errorf("failed to send heartbeat: %s", err.Error())
-		return err
+		return errors.Errorf("failed to send heartbeat: %s", err.Error())
 	}
 
+	// ttcode
+	// url := fmt.Sprintf("http://%s/plugin/topology/api/heartbeat?uuid=%s&addr=%s&interval=%d", conf.Config().Topo.Server_addr, m_u_struct.Agentuuid, addr, conf.Config().Topo.Heartbeat)
+	// resp, err := httputils.Get(url, nil)
+	// if err != nil {
+	// 	err = errors.Errorf("failed to send heartbeat: %s", err.Error())
+	// 	return err
+	// }
+
 	resp_body := &struct {
 		Code  int         `json:"code"`
 		Error string      `json:"error"`
 		Data  interface{} `json:"data"`
 	}{}
+
+	if len(resp.Body) == 0 {
+		return errors.New("heartbeat resp.body is nil")
+	}
+
 	err = json.Unmarshal(resp.Body, resp_body)
 	if err != nil {
-		err = errors.Errorf("failed to unmarshal json data: %s", err.Error())
-		return err
+		return errors.Errorf("failed to unmarshal json data: %s", err.Error())
 	}
 
 	if resp.StatusCode != http.StatusOK || resp_body.Code != 0 {
-		err = errors.Errorf("failed to send heartbeat: url => %s, statuscode => %d, code => %d", url, resp.StatusCode, resp_body.Code)
-		return err
+		return errors.Errorf("failed to send heartbeat: url => %s, statuscode => %d, code => %d", url, resp.StatusCode, resp_body.Code)
 	}
 
 	return nil
diff --git a/cmd/server/webserver/handle/basicHandler.go b/cmd/server/webserver/handle/basicHandler.go
index 0f09815..0c25c05 100644
--- a/cmd/server/webserver/handle/basicHandler.go
+++ b/cmd/server/webserver/handle/basicHandler.go
@@ -18,9 +18,11 @@ import (
 func HeartbeatHandle(ctx *gin.Context) {
 	// agent发送的心跳参数为uuid、ip:port、HeartbeatInterval、time，
 	// 写入redis的数据为 (heartbeat-<uuid>: {"UUID": "f7504bef-76e9-446c-95ee-196878b398a1", "Addr": "10.44.55.66:9992", "HeartbeatInterval": 60, "Time": "2023-12-22T17:09:23+08:00"})
+
 	value := redismanager.AgentHeartbeat{}
+
 	if err := ctx.ShouldBindJSON(&value); err != nil {
-		err := errors.New("bind json failed")
+		err := errors.Errorf("bind json failed: %s", err.Error())
 		ctx.JSON(http.StatusInternalServerError, gin.H{
 			"code":  -1,
 			"error": err.Error(),
@@ -28,6 +30,13 @@ func HeartbeatHandle(ctx *gin.Context) {
 		})
 		global.ERManager.ErrorTransmit("webserver", "error", err, true, true)
 	}
+	ctx.Request.Body.Close()
+
+	// ttcode
+	// value.UUID = ctx.Query("uuid")
+	// value.Addr = ctx.Query("addr")
+	// value.HeartbeatInterval, _ = strconv.Atoi(ctx.Query("interval"))
+
 	key := "heartbeat-topoagent-" + value.UUID
 	value.Time = time.Now()
 
@@ -43,7 +52,7 @@ func HeartbeatHandle(ctx *gin.Context) {
 	}
 
 	if agentmanager.Global_AgentManager.GetAgent_P(value.UUID) == nil {
-		err := errors.Errorf("unknown agent's heartbeat: %s, %s", value.UUID, value.Addr)
+		err := errors.Errorf("unknown agent's heartbeat: %s, %s, %+v", value.UUID, value.Addr, value)
 		global.ERManager.ErrorTransmit("webserver", "warn", err, false, false)
 		ctx.JSON(http.StatusUnauthorized, gin.H{
 			"code":  -1,
diff --git a/cmd/server/webserver/handle/customTopoHandler.go b/cmd/server/webserver/handle/customTopoHandler.go
index ca08a85..7f6a38b 100755
--- a/cmd/server/webserver/handle/customTopoHandler.go
+++ b/cmd/server/webserver/handle/customTopoHandler.go
@@ -160,7 +160,7 @@ func RunCustomTopoHandle(ctx *gin.Context) {
 				return
 			}
 			err = errors.Wrap(err, " ")
-			global.ERManager.ErrorTransmit("webserver", "error", err, false, true)
+			global.ERManager.ErrorTransmit("webserver", "error", err, false, false)
 			doneChan <- custom_topodata
 			response.Fail(ctx, nil, errors.Cause(err).Error())
 			return
-- 
Gitee