fix: NEZ-2021 调整更新流程,使用at 延迟1min 执行更新操作

3 files changed, 41 insertions, 22 deletions

diff --git a/src/main/java/net/geedge/confagent/controller/BaseController.java b/src/main/java/net/geedge/confagent/controller/BaseController.java
index 1ce7266..e9d4bef 100644
--- a/src/main/java/net/geedge/confagent/controller/BaseController.java
+++ b/src/main/java/net/geedge/confagent/controller/BaseController.java
@@ -15,6 +15,7 @@ public abstract class BaseController {
 	protected static final String PROMETHEUS_LISTEN_ADDR = "web.listen-address";
 	protected static final String SNMPEXPORTER_LISTEN_ADDR = "web.listen-address";
 	protected static final String BLACKBOXEXPORTER_LISTEN_ADDR = "web.listen-address";
+	protected static final String UPDATE_PACKAGE_PATH = "/tmp/nezha/nz-agent/%s/%s";
 
 	// loki push urls yml 文件 json 字符串
 	protected static String LOKI_PUSH_URL_JSONSTR;
diff --git a/src/main/java/net/geedge/confagent/controller/OTAController.java b/src/main/java/net/geedge/confagent/controller/OTAController.java
index 71636c5..ff8f8d2 100644
--- a/src/main/java/net/geedge/confagent/controller/OTAController.java
+++ b/src/main/java/net/geedge/confagent/controller/OTAController.java
@@ -1,6 +1,5 @@
 package net.geedge.confagent.controller;
 
-import cn.hutool.crypto.digest.MD5;
 import cn.hutool.log.Log;
 import net.geedge.confagent.util.*;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -11,7 +10,7 @@ import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.multipart.MultipartFile;
 
 import javax.servlet.http.HttpServletRequest;
-import java.io.File;
+import java.io.*;
 
 @RestController
 public class OTAController extends BaseController{
@@ -21,13 +20,13 @@ public class OTAController extends BaseController{
     @Autowired
     private ConfagentUtil confagentUtil;
 
+    private String updateSh = "/opt/nezha/nz-agent/update.sh";
+
     @Value("${version}")
     private String currentVersion;
 
-    private String rpmPackagePath = "/tmp/";
-
     @PostMapping("/ota")
-    public R ota(HttpServletRequest request, @RequestParam(name = "file", required = false) MultipartFile file, String md5){
+    public R ota(HttpServletRequest request, @RequestParam(name = "file", required = false) MultipartFile file, String md5 )throws IOException{
 
         //检查token
         String token = request.getHeader("Authorization");
@@ -37,18 +36,19 @@ public class OTAController extends BaseController{
             if (confagentUtil.checkToken(token).getCode() != RCode.SUCCESS.getCode())return R.error(RCode.AUTH_TOKEN_INVALID);
         }
 
-        //缓存文件
-        String filePath = rpmPackagePath + String.format("nezha.%s.rpm", CommonUtils.uuid());
-        log.info("cache rpm package name : {}", filePath);
-        File dest = new File(filePath);
+        //将文件缓存到 UPDATE_PACKAGE_PATH 目录下
+        String updateFileName = file.getOriginalFilename();
+        File updatePackagePath = Tool.FileUtil.file(String.format(UPDATE_PACKAGE_PATH , CommonUtils.uuid(), updateFileName));
+        updatePackagePath.mkdirs();
         try {
-            file.transferTo(dest);
+            file.transferTo(updatePackagePath);
         } catch (Exception e) {
             return R.error(RCode.OTA_FILE_CACHE_ERROR);
         }
 
-        //校验MD5
-        String updateFileMd5 = MD5.create().digestHex(dest);
+       //校验MD5
+        String updateFileMd5 = Tool.DigestUtil.md5Hex(new FileInputStream(updatePackagePath));
+        log.info("updateFileMd5 : {}",updateFileMd5);
         md5 = Tool.AesUtil.decrypt(md5, AES_SECRET_KEY);
         if(!updateFileMd5.equals(md5)){
             return R.error(RCode.OTA_MD5_ERROR);
@@ -63,16 +63,33 @@ public class OTAController extends BaseController{
         if( !(Tool.DateUtil.betweenDay(Tool.DateUtil.parseDate(updateVersion),Tool.DateUtil.parseDate(currentVersion), false) > 0)){
             return R.error(RCode.OTA_VERSION_ERROR);
         }
-        //更新
-        String updateCommand = String.format("rpm -Uvh %s", filePath);
-        log.info("update nz-agent : {}", updateCommand);
-        new Thread(() -> {
-            try {
-                Tool.RuntimeUtil.exec(updateCommand);
-            } catch (Exception e) {
-                log.error("update rpm package error : {}", e.getMessage());
-            }
-        }).start();
+
+        //创建update.sh文件 用于在线更新agent
+        String updateCommand = String.format("rpm -Uvh %s", updatePackagePath);
+        BufferedWriter out = null;
+        try {
+            out = new BufferedWriter(new FileWriter(updateSh));
+            out.write("#!/bin/bash");
+            out.newLine();
+            out.write(updateCommand);
+        } catch (IOException e) {
+            e.printStackTrace();
+        }finally {
+            out.close();
+        }
+
+        //检查atd服务是否启动
+        InputStream atdStatus = Tool.RuntimeUtil.exec("systemctl status atd").getInputStream();
+        String atdStatusStr = Tool.IoUtil.read(atdStatus).toString();
+        if(!atdStatusStr.contains("Active: active (running)")){
+            return R.error(RCode.OTA_ATD_SERVICE_DOWN);
+        }
+
+        //延迟一分钟更新
+        String atCommand ="at now +1 minute -f %s";
+        atCommand = String.format(atCommand,updateSh);
+        log.info("update nz-agent : {}", atCommand);
+        Tool.RuntimeUtil.exec(atCommand);
         return R.ok();
     }
 
diff --git a/src/main/java/net/geedge/confagent/util/RCode.java b/src/main/java/net/geedge/confagent/util/RCode.java
index 8e40b58..9e917c9 100644
--- a/src/main/java/net/geedge/confagent/util/RCode.java
+++ b/src/main/java/net/geedge/confagent/util/RCode.java
@@ -27,6 +27,7 @@ public enum RCode {
 	OTA_FILE_CACHE_ERROR(10114, "Cache rpm package failed"),
 	OTA_MD5_ERROR(10115, "MD5 is error"),
 	OTA_VERSION_ERROR(10116, "The updated version is lower than the current version"),
+	OTA_ATD_SERVICE_DOWN(10117, "The ATD service is down"),
 
     ERROR(999, "error"); //通用错误/未知错误
 	private RCode(Integer code, String msg) {