287 lines
7.8 KiB
Go
287 lines
7.8 KiB
Go
package httpapi
|
|
|
|
import (
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"mime"
|
|
"net/http"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"rk3588sys/agent/internal/files"
|
|
)
|
|
|
|
type agentBinaryUpdateResult struct {
|
|
StagingPath string `json:"staging_path"`
|
|
Sha256 string `json:"sha256"`
|
|
Size int64 `json:"size"`
|
|
MtimeMS int64 `json:"mtime_ms"`
|
|
ScriptPath string `json:"script_path"`
|
|
ExecPath string `json:"exec_path"`
|
|
Pid int `json:"pid"`
|
|
LogPath string `json:"log_path"`
|
|
TaskID string `json:"task_id"`
|
|
}
|
|
|
|
func (s *Server) handleAgentBinaryUpdate(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodPut {
|
|
errorJSON(w, http.StatusMethodNotAllowed, "method not allowed")
|
|
return
|
|
}
|
|
if !s.authorize(r, true) {
|
|
errorJSON(w, http.StatusUnauthorized, "unauthorized")
|
|
return
|
|
}
|
|
if mt, _, err := mime.ParseMediaType(r.Header.Get("Content-Type")); err != nil || mt != "application/octet-stream" {
|
|
errorJSON(w, http.StatusBadRequest, "validation failed: Content-Type must be application/octet-stream")
|
|
return
|
|
}
|
|
if r.ContentLength <= 0 {
|
|
errorJSON(w, http.StatusBadRequest, "validation failed: missing Content-Length")
|
|
return
|
|
}
|
|
if strings.TrimSpace(s.execPath) == "" {
|
|
errorJSON(w, http.StatusInternalServerError, "internal error: exec path is empty")
|
|
return
|
|
}
|
|
|
|
maxBytes := int64(s.agentCfg.MaxUploadMB) * 1024 * 1024
|
|
r.Body = http.MaxBytesReader(w, r.Body, maxBytes)
|
|
|
|
expected := strings.TrimSpace(r.Header.Get("X-Binary-Sha256"))
|
|
if expected != "" {
|
|
if len(expected) != 64 {
|
|
errorJSON(w, http.StatusBadRequest, "validation failed: invalid X-Binary-Sha256")
|
|
return
|
|
}
|
|
if _, err := hex.DecodeString(expected); err != nil {
|
|
errorJSON(w, http.StatusBadRequest, "validation failed: invalid X-Binary-Sha256")
|
|
return
|
|
}
|
|
}
|
|
|
|
baseDir := s.baseDir
|
|
if strings.TrimSpace(baseDir) == "" {
|
|
baseDir = filepath.Dir(s.execPath)
|
|
}
|
|
updateDir := filepath.Join(baseDir, "updates", "agent")
|
|
stagingDir := filepath.Join(updateDir, "staging")
|
|
backupDir := filepath.Join(updateDir, "backup")
|
|
lockPath := filepath.Join(updateDir, "update.lock")
|
|
if err := files.EnsureDir(updateDir, 0o755); err != nil {
|
|
errorJSON(w, http.StatusInternalServerError, "internal error: "+err.Error())
|
|
return
|
|
}
|
|
|
|
lockFile, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o644)
|
|
if err != nil {
|
|
if os.IsExist(err) {
|
|
errorJSON(w, http.StatusConflict, "update already in progress")
|
|
return
|
|
}
|
|
errorJSON(w, http.StatusInternalServerError, "internal error: create lock failed: "+err.Error())
|
|
return
|
|
}
|
|
_, _ = lockFile.WriteString(strconv.Itoa(os.Getpid()) + "\n")
|
|
_ = lockFile.Close()
|
|
lockOK := false
|
|
defer func() {
|
|
if !lockOK {
|
|
_ = os.Remove(lockPath)
|
|
}
|
|
}()
|
|
|
|
task := s.tasks.Start("agent_binary_update")
|
|
stagingPath, sha, size, mtimeMS, err := writeStagingBinary(stagingDir, r.Body, r.ContentLength, expected)
|
|
if err != nil {
|
|
_, _ = s.tasks.Finish(task.ID, nil, err)
|
|
s.recordAudit(r, "agent.binary.update", false, err.Error())
|
|
if strings.Contains(err.Error(), "payload too large") || strings.Contains(err.Error(), "request body too large") {
|
|
errorJSON(w, http.StatusRequestEntityTooLarge, "payload too large")
|
|
return
|
|
}
|
|
if strings.Contains(err.Error(), "sha256 mismatch") {
|
|
errorJSON(w, http.StatusBadRequest, "validation failed: sha256 mismatch")
|
|
return
|
|
}
|
|
errorJSON(w, http.StatusInternalServerError, "internal error: "+err.Error())
|
|
return
|
|
}
|
|
|
|
pid := os.Getpid()
|
|
logPath := filepath.Join(updateDir, "update.log")
|
|
scriptPath, err := writeAgentUpdateScript(updateDir, pid, s.execPath, s.agentCfg.ConfigPath, stagingPath, backupDir, logPath, lockPath)
|
|
if err != nil {
|
|
_, _ = s.tasks.Finish(task.ID, nil, err)
|
|
s.recordAudit(r, "agent.binary.update", false, err.Error())
|
|
errorJSON(w, http.StatusInternalServerError, "internal error: "+err.Error())
|
|
return
|
|
}
|
|
|
|
cmd := exec.Command("/bin/sh", scriptPath)
|
|
if err := cmd.Start(); err != nil {
|
|
_, _ = s.tasks.Finish(task.ID, nil, err)
|
|
s.recordAudit(r, "agent.binary.update", false, err.Error())
|
|
errorJSON(w, http.StatusInternalServerError, "internal error: start update script failed: "+err.Error())
|
|
return
|
|
}
|
|
|
|
res := agentBinaryUpdateResult{
|
|
StagingPath: filepath.ToSlash(stagingPath),
|
|
Sha256: sha,
|
|
Size: size,
|
|
MtimeMS: mtimeMS,
|
|
ScriptPath: filepath.ToSlash(scriptPath),
|
|
ExecPath: filepath.ToSlash(s.execPath),
|
|
Pid: pid,
|
|
LogPath: filepath.ToSlash(logPath),
|
|
TaskID: task.ID,
|
|
}
|
|
_, _ = s.tasks.Finish(task.ID, res, nil)
|
|
lockOK = true
|
|
s.recordAudit(r, "agent.binary.update", true, res.StagingPath)
|
|
writeJSON(w, http.StatusOK, map[string]any{"ok": true, "update": res})
|
|
|
|
go func() {
|
|
time.Sleep(500 * time.Millisecond)
|
|
os.Exit(0)
|
|
}()
|
|
}
|
|
|
|
func writeStagingBinary(dir string, r io.Reader, contentLength int64, expectedSha256 string) (string, string, int64, int64, error) {
|
|
if err := files.EnsureDir(dir, 0o755); err != nil {
|
|
return "", "", 0, 0, err
|
|
}
|
|
f, err := os.CreateTemp(dir, ".tmp-*")
|
|
if err != nil {
|
|
return "", "", 0, 0, fmt.Errorf("create temp: %w", err)
|
|
}
|
|
tmp := f.Name()
|
|
ok := false
|
|
defer func() {
|
|
_ = f.Close()
|
|
if !ok {
|
|
_ = os.Remove(tmp)
|
|
}
|
|
}()
|
|
|
|
h := sha256.New()
|
|
mw := io.MultiWriter(f, h)
|
|
if _, err := io.CopyN(mw, r, contentLength); err != nil {
|
|
return "", "", 0, 0, fmt.Errorf("read body: %w", err)
|
|
}
|
|
if err := f.Chmod(0o755); err != nil {
|
|
return "", "", 0, 0, fmt.Errorf("chmod temp: %w", err)
|
|
}
|
|
if err := f.Sync(); err != nil {
|
|
return "", "", 0, 0, fmt.Errorf("fsync temp: %w", err)
|
|
}
|
|
if err := f.Close(); err != nil {
|
|
return "", "", 0, 0, fmt.Errorf("close temp: %w", err)
|
|
}
|
|
|
|
sha := hex.EncodeToString(h.Sum(nil))
|
|
if expectedSha256 != "" && !strings.EqualFold(expectedSha256, sha) {
|
|
return "", "", 0, 0, errors.New("sha256 mismatch")
|
|
}
|
|
|
|
stagingPath := filepath.Join(dir, "agent_new.bin")
|
|
if err := files.ReplaceFile(tmp, stagingPath); err != nil {
|
|
return "", "", 0, 0, err
|
|
}
|
|
st, err := os.Stat(stagingPath)
|
|
if err != nil {
|
|
return "", "", 0, 0, fmt.Errorf("stat staging: %w", err)
|
|
}
|
|
ok = true
|
|
return stagingPath, sha, st.Size(), st.ModTime().UnixMilli(), nil
|
|
}
|
|
|
|
func writeAgentUpdateScript(updateDir string, pid int, execPath, configPath, stagingPath, backupDir, logPath, lockPath string) (string, error) {
|
|
if err := files.EnsureDir(updateDir, 0o755); err != nil {
|
|
return "", err
|
|
}
|
|
if err := files.EnsureDir(backupDir, 0o755); err != nil {
|
|
return "", err
|
|
}
|
|
scriptPath := filepath.Join(updateDir, "apply_update.sh")
|
|
|
|
content := fmt.Sprintf(`#!/bin/sh
|
|
set -u
|
|
|
|
PID=%d
|
|
EXEC=%s
|
|
CONFIG=%s
|
|
STAGING=%s
|
|
BACKUP_DIR=%s
|
|
LOG=%s
|
|
LOCK=%s
|
|
|
|
trap 'rm -f "$LOCK"' EXIT
|
|
|
|
mkdir -p "$BACKUP_DIR"
|
|
echo "$(date -Iseconds) update start" >> "$LOG"
|
|
|
|
if kill -0 "$PID" 2>/dev/null; then
|
|
i=0
|
|
while kill -0 "$PID" 2>/dev/null; do
|
|
i=$((i+1))
|
|
if [ "$i" -ge 60 ]; then
|
|
break
|
|
fi
|
|
sleep 0.5
|
|
done
|
|
fi
|
|
if kill -0 "$PID" 2>/dev/null; then
|
|
kill -TERM "$PID" 2>/dev/null || true
|
|
sleep 2
|
|
fi
|
|
if kill -0 "$PID" 2>/dev/null; then
|
|
kill -KILL "$PID" 2>/dev/null || true
|
|
sleep 1
|
|
fi
|
|
if kill -0 "$PID" 2>/dev/null; then
|
|
echo "$(date -Iseconds) agent still running" >> "$LOG"
|
|
exit 1
|
|
fi
|
|
|
|
ts=$(date +%%Y%%m%%d-%%H%%M%%S)
|
|
backup="$BACKUP_DIR/agent.bak.$ts"
|
|
if [ -f "$EXEC" ]; then
|
|
mv "$EXEC" "$backup"
|
|
fi
|
|
|
|
if mv "$STAGING" "$EXEC"; then
|
|
chmod 755 "$EXEC"
|
|
else
|
|
if [ -f "$backup" ]; then
|
|
mv "$backup" "$EXEC"
|
|
fi
|
|
echo "$(date -Iseconds) replace failed" >> "$LOG"
|
|
exit 1
|
|
fi
|
|
|
|
nohup "$EXEC" --config "$CONFIG" >/dev/null 2>&1 &
|
|
echo "$(date -Iseconds) update done" >> "$LOG"
|
|
`, pid, shQuote(execPath), shQuote(configPath), shQuote(stagingPath), shQuote(backupDir), shQuote(logPath), shQuote(lockPath))
|
|
|
|
if err := files.WriteFileAtomic(scriptPath, []byte(content), 0o755); err != nil {
|
|
return "", err
|
|
}
|
|
return scriptPath, nil
|
|
}
|
|
|
|
func shQuote(s string) string {
|
|
if s == "" {
|
|
return "''"
|
|
}
|
|
return "'" + strings.ReplaceAll(s, "'", `'"'"'`) + "'"
|
|
}
|