parse container resp

nm/nm_operate.go

@@ -73,6 +73,7 @@ func inputNodeManagerChan(manager *NodeManager, nodeManagerClient *models.NodeMa
 			IsSelected:    isSelect,
 			LastHeartTime: time.Now(),
 		}
+		usedNodeManagerClient = append(usedNodeManagerClient, nodeManagerClient)
 	}
 	serviceClient := operate.ConnNmGrpc(manager.Info.Endpoint)
 	if serviceClient == nil {
@@ -82,7 +83,6 @@ func inputNodeManagerChan(manager *NodeManager, nodeManagerClient *models.NodeMa
 	nodeManagerClient.Status = true
 	nodeManagerClient.Client = serviceClient
 	nodeManagerClientChan <- nodeManagerClient
-	usedNodeManagerClient = append(usedNodeManagerClient, nodeManagerClient)
 	manager.IsUsed = true
 	return true
 }

nm/task_msg.go

@@ -2,7 +2,6 @@ package nm
 
 import (
 	"bytes"
-	"encoding/base64"
 	"encoding/json"
 	"example.com/m/conf"
 	"example.com/m/log"
@@ -272,20 +271,22 @@ func (t *TaskHandler) ComputeTaskHandler(taskMsg *nodeManagerV1.PushTaskMessage)
 					default:
 						log.Error("data is unknown type", v)
 					}
-					if base64ImageStr != "" && utils.IsBase64ImageStr(base64ImageStr) {
-						containerRespOutput := strings.SplitN(base64ImageStr, ",", 2)
-						imageStr := containerRespOutput[1]
-						queryString := utils.MatchFileCacheQueryString(taskParam.Headers, taskCmd.ImageName, t.DockerOp.ModelsInfo, containerRespOutput[0])
-						prefix := strings.Split(strings.Split(containerRespOutput[0], ";")[0], ":")[1]
-						ossUri, err := t.uploadOSS(taskMsg.TaskId, queryString, imageStr, prefix)
-						if err != nil {
-							log.WithError(err).Error("upload image into file cache failed")
-							return
-						}
-						log.WithField("uri", ossUri).Info("upload image OSS  successful")
-						if ossUri != "" {
-							taskExecResult.TaskHttpStatusCode = models.RedirectCode
-							post.Header.Set("Location", ossUri)
+					if base64ImageStr != "" {
+						isBase64, decodeByte, respFormat, suffix := utils.IsBase64ImageStr(base64ImageStr)
+						if isBase64 {
+							queryString := utils.MatchFileCacheQueryString(taskParam.Headers, taskCmd.ImageName, t.DockerOp.ModelsInfo, respFormat)
+							ossUri, err := t.uploadOSS(taskMsg.TaskId, queryString, decodeByte, suffix)
+							if err != nil {
+								log.WithError(err).Error("upload image into file cache failed")
+								return
+							}
+							log.WithField("uri", ossUri).Info("upload image OSS  successful")
+							if ossUri != "" {
+								taskExecResult.TaskHttpStatusCode = models.RedirectCode
+								post.Header.Set("Location", ossUri)
+							}
+						} else {
+							taskExecResult.TaskRespBody = respBody
 						}
 					} else {
 						taskExecResult.TaskRespBody = respBody
@@ -371,21 +372,15 @@ func (t *TaskHandler) foundImageIsRunning(imageId string) (bool, string, uint16)
 	return false, "", 0
 }
 
-func (t *TaskHandler) uploadOSS(taskId string, queries string, base64Image, prefix string) (string, error) {
+func (t *TaskHandler) uploadOSS(taskId string, queries string, decodedImage []byte, suffix string) (string, error) {
 	var requestBody bytes.Buffer
 	writer := multipart.NewWriter(&requestBody)
 	// 创建文件表单字段
-	fileField, err := writer.CreateFormFile("file", fmt.Sprintf("%s.%s", taskId, prefix))
+	fileField, err := writer.CreateFormFile("file", fmt.Sprintf("%s.%s", taskId, suffix))
 	if err != nil {
 		log.WithError(err).Error("Error creating form file")
 		return "", err
 	}
-	// 将 base64 解码后的内容复制到表单字段
-	decodedImage, err := base64.StdEncoding.DecodeString(base64Image)
-	if err != nil {
-		log.WithError(err).Error("Error decoding base64 image")
-		return "", err
-	}
 	_, err = io.Copy(fileField, bytes.NewReader(decodedImage))
 	//_, err = io.Copy(fileField, strings.NewReader(base64Image))
 	if err != nil {

utils/util.go

@@ -96,14 +96,28 @@ func readPrivateKey() (*ecdsa.PrivateKey, error) {
 }
 
 // IsBase64ImageStr 检查字符串是否是 Base64 编码的图像数据
-func IsBase64ImageStr(imageStr string) bool {
+func IsBase64ImageStr(imageStr string) (bool, []byte, string, string) {
 	// 移除可能的前缀（如 "data:image/png;base64,"）
-	imageStr = strings.SplitN(imageStr, ",", 2)[1]
-	_, err := base64.StdEncoding.DecodeString(imageStr)
-	return err == nil
+	imageStrArr := strings.SplitN(imageStr, ",", 2)
+	base64CodePrefix := ""
+	base64Code := ""
+	if len(imageStrArr) == 1 {
+		base64Code = imageStrArr[0]
+		base64CodePrefix = "data:image/png;base64"
+	} else {
+		base64CodePrefix = imageStrArr[0]
+		base64Code = imageStrArr[1]
+	}
+	decodeBytes, err := base64.StdEncoding.DecodeString(base64Code)
+	if err != nil {
+		return false, nil, "", ""
+	}
+	formatStr := strings.Split(strings.Split(base64CodePrefix, ";")[0], ":")[1]
+	suffix := strings.Split(formatStr, "/")[1]
+	return true, decodeBytes, formatStr, suffix
 }
 
-func MatchFileCacheQueryString(params map[string][]string, taskImageName string, modelsInfo []*models.ModelInfo, formatType string) string {
+func MatchFileCacheQueryString(params map[string][]string, taskImageName string, modelsInfo []*models.ModelInfo, contentType string) string {
 	values := url.Values{}
 	isExistFileExpires := false
 	for key, value := range params {
@@ -126,7 +140,6 @@ func MatchFileCacheQueryString(params map[string][]string, taskImageName string,
 	if !isModelExistFileExpires {
 		values.Add(models.ResultFileExpiresDB, "600")
 	}
-	contentType := strings.Split(strings.Split(formatType, ";")[0], ":")[1]
 	values.Add(models.ContentType, contentType)
 	return values.Encode()
 }