AI-tianyan/internal/control/ota_agent.go

package control

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"tianyan-edge/internal/config"
)

type OTAAgent struct {
	cfg *config.Config
}

func NewOTAAgent(cfg *config.Config) *OTAAgent {
	return &OTAAgent{cfg: cfg}
}

func (o *OTAAgent) Run(ctx context.Context) {
	client := &http.Client{Timeout: 30 * time.Second}
	ticker := time.NewTicker(10 * time.Minute)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			if inMaintenanceWindow() {
				o.check(client)
			}
		}
	}
}

func inMaintenanceWindow() bool {
	h := time.Now().Hour()
	return h >= 22 || h < 6
}

func (o *OTAAgent) check(client *http.Client) {
	url := fmt.Sprintf("%s/api/v1/edge/update/manifest?edge_id=%s&current_version=%s",
		o.cfg.CloudURL, o.cfg.EdgeID, o.cfg.Version)
	req, _ := http.NewRequest(http.MethodGet, url, nil)
	req.Header.Set("Authorization", "Bearer "+o.cfg.EdgeToken)
	resp, err := client.Do(req)
	if err != nil {
		log.Printf("ota manifest: %v", err)
		return
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return
	}
	var mf map[string]string
	if err := json.NewDecoder(resp.Body).Decode(&mf); err != nil {
		return
	}
	target := mf["target_version"]
	if target == "" || target == o.cfg.Version {
		return
	}
	if err := o.downloadAndVerify(client, mf); err != nil {
		o.report(client, target, "failed")
		return
	}
	o.report(client, target, "verified")
}

func (o *OTAAgent) downloadAndVerify(client *http.Client, mf map[string]string) error {
	url, target := mf["package_url"], mf["target_version"]
	expected := mf["sha256"]
	if url == "" || expected == "" {
		return fmt.Errorf("invalid manifest")
	}
	resp, err := client.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	sum := sha256.Sum256(data)
	actual := hex.EncodeToString(sum[:])
	if actual != expected {
		return fmt.Errorf("checksum mismatch")
	}
	_ = os.MkdirAll("/opt/tianyan-edge/staging", 0o755)
	path := filepath.Join("/opt/tianyan-edge/staging", target+".tar.gz")
	return os.WriteFile(path, data, 0o644)
}

func (o *OTAAgent) report(client *http.Client, version, status string) {
	payload := map[string]any{"edge_id": o.cfg.EdgeID, "version": version,
		"status": status, "ts": float64(time.Now().UnixMilli()) / 1000.0}
	body, _ := json.Marshal(payload)
	url := fmt.Sprintf("%s/api/v1/edge/update/report", o.cfg.CloudURL)
	req, _ := http.NewRequest(http.MethodPost, url, bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+o.cfg.EdgeToken)
	resp, err := client.Do(req)
	if err == nil {
		resp.Body.Close()
	}
}