1120 lines
27 KiB
Go
1120 lines
27 KiB
Go
package main
|
||
|
||
import (
|
||
"archive/zip"
|
||
"bufio"
|
||
"context"
|
||
"crypto/rand"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"os"
|
||
"os/exec"
|
||
"path/filepath"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/gorilla/websocket"
|
||
"golang.org/x/sync/errgroup"
|
||
)
|
||
|
||
// 任务状态
|
||
type TaskStatus string
|
||
|
||
const (
|
||
StatusPending TaskStatus = "pending"
|
||
StatusRunning TaskStatus = "running"
|
||
StatusCompleted TaskStatus = "completed"
|
||
StatusFailed TaskStatus = "failed"
|
||
)
|
||
|
||
// 镜像下载任务
|
||
type ImageTask struct {
|
||
Image string `json:"image"`
|
||
Progress float64 `json:"progress"`
|
||
Status string `json:"status"`
|
||
Error string `json:"error,omitempty"`
|
||
OutputPath string `json:"-"` // 输出文件路径，不发送给客户端
|
||
lock sync.Mutex `json:"-"` // 镜像任务自己的锁
|
||
}
|
||
|
||
// 下载任务
|
||
type DownloadTask struct {
|
||
ID string `json:"id"`
|
||
Images []*ImageTask `json:"images"`
|
||
CompletedCount int `json:"completedCount"` // 已完成任务数
|
||
TotalCount int `json:"totalCount"` // 总任务数
|
||
Status TaskStatus `json:"status"`
|
||
OutputFile string `json:"-"` // 最终输出文件
|
||
TempDir string `json:"-"` // 临时目录
|
||
StatusLock sync.RWMutex `json:"-"` // 状态锁，使用读写锁提高并发性
|
||
ProgressLock sync.RWMutex `json:"-"` // 进度锁
|
||
ImageLock sync.RWMutex `json:"-"` // 镜像列表锁
|
||
updateChan chan *ProgressUpdate `json:"-"` // 进度更新通道
|
||
}
|
||
|
||
// 进度更新消息
|
||
type ProgressUpdate struct {
|
||
TaskID string
|
||
ImageIndex int
|
||
Progress float64
|
||
Status string
|
||
Error string
|
||
}
|
||
|
||
// WebSocket客户端
|
||
type Client struct {
|
||
Conn *websocket.Conn
|
||
TaskID string
|
||
Send chan []byte
|
||
CloseOnce sync.Once
|
||
}
|
||
|
||
var (
|
||
// WebSocket升级器
|
||
upgrader = websocket.Upgrader{
|
||
ReadBufferSize: 1024,
|
||
WriteBufferSize: 1024,
|
||
CheckOrigin: func(r *http.Request) bool {
|
||
return true // 允许所有源
|
||
},
|
||
}
|
||
|
||
// 活跃任务映射
|
||
tasks = make(map[string]*DownloadTask)
|
||
tasksLock sync.Mutex
|
||
clients = make(map[string]*Client)
|
||
clientLock sync.Mutex
|
||
)
|
||
|
||
// 初始化Skopeo相关路由
|
||
func initSkopeoRoutes(router *gin.Engine) {
|
||
// 创建临时目录
|
||
os.MkdirAll("./temp", 0755)
|
||
|
||
// WebSocket路由 - 用于实时获取进度
|
||
router.GET("/ws/:taskId", handleWebSocket)
|
||
|
||
// 创建下载任务，应用限流中间件
|
||
ApplyRateLimit(router, "/api/download", "POST", handleDownload)
|
||
|
||
// 获取任务状态
|
||
router.GET("/api/task/:taskId", getTaskStatus)
|
||
|
||
// 下载文件
|
||
router.GET("/api/files/:filename", serveFile)
|
||
|
||
// 启动清理过期文件的goroutine
|
||
go cleanupTempFiles()
|
||
}
|
||
|
||
// 处理WebSocket连接
|
||
func handleWebSocket(c *gin.Context) {
|
||
taskID := c.Param("taskId")
|
||
|
||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||
if err != nil {
|
||
fmt.Printf("WebSocket升级失败: %v\n", err)
|
||
return
|
||
}
|
||
|
||
client := &Client{
|
||
Conn: conn,
|
||
TaskID: taskID,
|
||
Send: make(chan []byte, 256),
|
||
}
|
||
|
||
// 注册客户端
|
||
clientLock.Lock()
|
||
clients[taskID] = client
|
||
clientLock.Unlock()
|
||
|
||
// 启动goroutine处理消息发送
|
||
go client.writePump()
|
||
|
||
// 如果任务已存在，立即发送当前状态
|
||
tasksLock.Lock()
|
||
if task, exists := tasks[taskID]; exists {
|
||
tasksLock.Unlock()
|
||
taskJSON, _ := json.Marshal(task)
|
||
client.Send <- taskJSON
|
||
} else {
|
||
tasksLock.Unlock()
|
||
}
|
||
|
||
// 处理WebSocket关闭
|
||
conn.SetCloseHandler(func(code int, text string) error {
|
||
client.CloseOnce.Do(func() {
|
||
close(client.Send)
|
||
clientLock.Lock()
|
||
delete(clients, taskID)
|
||
clientLock.Unlock()
|
||
})
|
||
return nil
|
||
})
|
||
}
|
||
|
||
// 客户端消息发送loop
|
||
func (c *Client) writePump() {
|
||
defer func() {
|
||
c.Conn.Close()
|
||
}()
|
||
|
||
for message := range c.Send {
|
||
err := c.Conn.WriteMessage(websocket.TextMessage, message)
|
||
if err != nil {
|
||
fmt.Printf("发送WS消息失败: %v\n", err)
|
||
break
|
||
}
|
||
}
|
||
}
|
||
|
||
// 获取任务状态
|
||
func getTaskStatus(c *gin.Context) {
|
||
taskID := c.Param("taskId")
|
||
|
||
tasksLock.Lock()
|
||
task, exists := tasks[taskID]
|
||
tasksLock.Unlock()
|
||
|
||
if !exists {
|
||
c.JSON(http.StatusNotFound, gin.H{"error": "任务不存在"})
|
||
return
|
||
}
|
||
|
||
// 创建任务状态副本以避免序列化过程中的锁
|
||
taskCopy := &DownloadTask{
|
||
ID: task.ID,
|
||
CompletedCount: 0,
|
||
TotalCount: len(task.Images),
|
||
Status: TaskStatus(""),
|
||
Images: nil,
|
||
}
|
||
|
||
// 复制状态信息
|
||
task.StatusLock.RLock()
|
||
taskCopy.Status = task.Status
|
||
task.StatusLock.RUnlock()
|
||
|
||
task.ProgressLock.RLock()
|
||
taskCopy.CompletedCount = task.CompletedCount
|
||
task.ProgressLock.RUnlock()
|
||
|
||
// 复制镜像信息
|
||
task.ImageLock.RLock()
|
||
taskCopy.Images = make([]*ImageTask, len(task.Images))
|
||
for i, img := range task.Images {
|
||
img.lock.Lock()
|
||
taskCopy.Images[i] = &ImageTask{
|
||
Image: img.Image,
|
||
Progress: img.Progress,
|
||
Status: img.Status,
|
||
Error: img.Error,
|
||
}
|
||
img.lock.Unlock()
|
||
}
|
||
task.ImageLock.RUnlock()
|
||
|
||
c.JSON(http.StatusOK, taskCopy)
|
||
}
|
||
|
||
// 生成随机任务ID
|
||
func generateTaskID() string {
|
||
b := make([]byte, 16)
|
||
rand.Read(b)
|
||
return hex.EncodeToString(b)
|
||
}
|
||
|
||
// 初始化任务并启动进度处理器
|
||
func initTask(task *DownloadTask) {
|
||
// 创建进度更新通道
|
||
task.updateChan = make(chan *ProgressUpdate, 100)
|
||
|
||
// 启动进度处理goroutine
|
||
go func() {
|
||
for update := range task.updateChan {
|
||
if update == nil {
|
||
// 通道关闭信号
|
||
break
|
||
}
|
||
|
||
// 获取更新的镜像
|
||
task.ImageLock.RLock()
|
||
if update.ImageIndex < 0 || update.ImageIndex >= len(task.Images) {
|
||
task.ImageLock.RUnlock()
|
||
continue
|
||
}
|
||
imgTask := task.Images[update.ImageIndex]
|
||
task.ImageLock.RUnlock()
|
||
|
||
statusChanged := false
|
||
prevStatus := ""
|
||
|
||
// 更新镜像进度和状态
|
||
imgTask.lock.Lock()
|
||
if update.Progress > 0 {
|
||
imgTask.Progress = update.Progress
|
||
}
|
||
if update.Status != "" && update.Status != imgTask.Status {
|
||
prevStatus = imgTask.Status
|
||
imgTask.Status = update.Status
|
||
statusChanged = true
|
||
}
|
||
if update.Error != "" {
|
||
imgTask.Error = update.Error
|
||
}
|
||
imgTask.lock.Unlock()
|
||
|
||
// 检查状态变化并更新完成计数
|
||
if statusChanged {
|
||
task.ProgressLock.Lock()
|
||
// 如果之前不是Completed，现在是Completed，增加计数
|
||
if prevStatus != string(StatusCompleted) && update.Status == string(StatusCompleted) {
|
||
task.CompletedCount++
|
||
fmt.Printf("任务 %s: 镜像 %d 完成，当前完成数: %d/%d\n",
|
||
task.ID, update.ImageIndex, task.CompletedCount, task.TotalCount)
|
||
}
|
||
// 如果之前是Completed，现在不是，减少计数
|
||
if prevStatus == string(StatusCompleted) && update.Status != string(StatusCompleted) {
|
||
task.CompletedCount--
|
||
if task.CompletedCount < 0 {
|
||
task.CompletedCount = 0
|
||
}
|
||
}
|
||
task.ProgressLock.Unlock()
|
||
}
|
||
|
||
// 发送更新到客户端
|
||
sendTaskUpdate(task)
|
||
}
|
||
}()
|
||
}
|
||
|
||
// 发送进度更新
|
||
func sendProgressUpdate(task *DownloadTask, index int, progress float64, status string, errorMsg string) {
|
||
select {
|
||
case task.updateChan <- &ProgressUpdate{
|
||
TaskID: task.ID,
|
||
ImageIndex: index,
|
||
Progress: progress,
|
||
Status: status,
|
||
Error: errorMsg,
|
||
}:
|
||
// 成功发送
|
||
default:
|
||
// 通道已满，丢弃更新
|
||
fmt.Printf("Warning: Update channel for task %s is full\n", task.ID)
|
||
}
|
||
}
|
||
|
||
// 更新总进度 - 重新计算已完成任务数
|
||
func updateTaskTotalProgress(task *DownloadTask) {
|
||
task.ProgressLock.Lock()
|
||
defer task.ProgressLock.Unlock()
|
||
|
||
completedCount := 0
|
||
|
||
task.ImageLock.RLock()
|
||
totalCount := len(task.Images)
|
||
task.ImageLock.RUnlock()
|
||
|
||
if totalCount == 0 {
|
||
return
|
||
}
|
||
|
||
task.ImageLock.RLock()
|
||
for _, img := range task.Images {
|
||
img.lock.Lock()
|
||
if img.Status == string(StatusCompleted) {
|
||
completedCount++
|
||
}
|
||
img.lock.Unlock()
|
||
}
|
||
task.ImageLock.RUnlock()
|
||
|
||
task.CompletedCount = completedCount
|
||
task.TotalCount = totalCount
|
||
|
||
fmt.Printf("任务 %s: 进度更新 %d/%d 已完成\n", task.ID, completedCount, totalCount)
|
||
}
|
||
|
||
// 处理下载请求
|
||
func handleDownload(c *gin.Context) {
|
||
type DownloadRequest struct {
|
||
Images []string `json:"images"`
|
||
Platform string `json:"platform"` // 平台: amd64, arm64等
|
||
}
|
||
|
||
var req DownloadRequest
|
||
if err := c.ShouldBindJSON(&req); err != nil {
|
||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数"})
|
||
return
|
||
}
|
||
|
||
if len(req.Images) == 0 {
|
||
c.JSON(http.StatusBadRequest, gin.H{"error": "请提供至少一个镜像"})
|
||
return
|
||
}
|
||
|
||
// 添加镜像数量限制10个，防止恶意刷流量
|
||
if len(req.Images) > 10 {
|
||
c.JSON(http.StatusBadRequest, gin.H{"error": "您下载的数量太多，宝宝承受不住"})
|
||
return
|
||
}
|
||
|
||
// 创建新任务
|
||
taskID := generateTaskID()
|
||
tempDir := filepath.Join("./temp", taskID)
|
||
os.MkdirAll(tempDir, 0755)
|
||
|
||
// 初始化任务
|
||
imageTasks := make([]*ImageTask, len(req.Images))
|
||
for i, image := range req.Images {
|
||
imageTasks[i] = &ImageTask{
|
||
Image: image,
|
||
Progress: 0,
|
||
Status: string(StatusPending),
|
||
}
|
||
}
|
||
|
||
task := &DownloadTask{
|
||
ID: taskID,
|
||
Images: imageTasks,
|
||
CompletedCount: 0,
|
||
TotalCount: len(imageTasks),
|
||
Status: StatusPending,
|
||
TempDir: tempDir,
|
||
}
|
||
|
||
// 初始化任务通道和处理器
|
||
initTask(task)
|
||
|
||
// 保存任务
|
||
tasksLock.Lock()
|
||
tasks[taskID] = task
|
||
tasksLock.Unlock()
|
||
|
||
// 异步处理下载
|
||
go func() {
|
||
processDownloadTask(task, req.Platform)
|
||
// 任务完成后关闭更新通道
|
||
close(task.updateChan)
|
||
}()
|
||
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"taskId": taskID,
|
||
"status": "started",
|
||
"totalCount": len(imageTasks),
|
||
})
|
||
}
|
||
|
||
// 处理下载任务
|
||
func processDownloadTask(task *DownloadTask, platform string) {
|
||
// 设置任务状态为运行中
|
||
task.StatusLock.Lock()
|
||
task.Status = StatusRunning
|
||
task.StatusLock.Unlock()
|
||
|
||
// 初始化总任务数
|
||
task.ImageLock.RLock()
|
||
imageCount := len(task.Images)
|
||
task.ImageLock.RUnlock()
|
||
|
||
task.ProgressLock.Lock()
|
||
task.TotalCount = imageCount
|
||
task.CompletedCount = 0
|
||
task.ProgressLock.Unlock()
|
||
|
||
// 通知客户端任务已开始
|
||
sendTaskUpdate(task)
|
||
|
||
// 创建错误组，用于管理所有下载goroutine
|
||
ctx, cancel := context.WithCancel(context.Background())
|
||
defer cancel() // 确保资源被释放
|
||
|
||
g, ctx := errgroup.WithContext(ctx)
|
||
|
||
// 启动并发下载
|
||
task.ImageLock.RLock()
|
||
imageCount = len(task.Images)
|
||
task.ImageLock.RUnlock()
|
||
|
||
// 创建工作池限制并发数
|
||
const maxConcurrent = 5
|
||
semaphore := make(chan struct{}, maxConcurrent)
|
||
|
||
// 添加下载任务
|
||
for i := 0; i < imageCount; i++ {
|
||
index := i // 捕获循环变量
|
||
|
||
g.Go(func() error {
|
||
// 获取信号量，限制并发
|
||
semaphore <- struct{}{}
|
||
defer func() { <-semaphore }()
|
||
|
||
task.ImageLock.RLock()
|
||
imgTask := task.Images[index]
|
||
task.ImageLock.RUnlock()
|
||
|
||
// 下载镜像
|
||
err := downloadImageWithContext(ctx, task, index, imgTask, platform)
|
||
if err != nil {
|
||
fmt.Printf("镜像 %s 下载失败: %v\n", imgTask.Image, err)
|
||
return err
|
||
}
|
||
return nil
|
||
})
|
||
}
|
||
|
||
// 等待所有下载完成
|
||
err := g.Wait()
|
||
|
||
// 再次计算已完成任务数，确保正确
|
||
updateTaskTotalProgress(task)
|
||
|
||
// 检查是否有错误发生
|
||
if err != nil {
|
||
task.StatusLock.Lock()
|
||
task.Status = StatusFailed
|
||
task.StatusLock.Unlock()
|
||
sendTaskUpdate(task)
|
||
return
|
||
}
|
||
|
||
// 判断是单个tar还是需要打包
|
||
var finalFilePath string
|
||
|
||
task.StatusLock.Lock()
|
||
|
||
// 检查是否所有镜像都下载成功
|
||
allSuccess := true
|
||
task.ImageLock.RLock()
|
||
for _, img := range task.Images {
|
||
img.lock.Lock()
|
||
if img.Status != string(StatusCompleted) {
|
||
allSuccess = false
|
||
}
|
||
img.lock.Unlock()
|
||
}
|
||
task.ImageLock.RUnlock()
|
||
|
||
if !allSuccess {
|
||
task.Status = StatusFailed
|
||
task.StatusLock.Unlock()
|
||
sendTaskUpdate(task)
|
||
return
|
||
}
|
||
|
||
// 如果只有一个文件，直接使用它
|
||
task.ImageLock.RLock()
|
||
if imageCount == 1 {
|
||
imgTask := task.Images[0]
|
||
imgTask.lock.Lock()
|
||
if imgTask.Status == string(StatusCompleted) {
|
||
finalFilePath = imgTask.OutputPath
|
||
// 重命名为更友好的名称
|
||
imageName := strings.ReplaceAll(imgTask.Image, "/", "_")
|
||
imageName = strings.ReplaceAll(imageName, ":", "_")
|
||
newPath := filepath.Join(task.TempDir, imageName+".tar")
|
||
os.Rename(finalFilePath, newPath)
|
||
finalFilePath = newPath
|
||
}
|
||
imgTask.lock.Unlock()
|
||
} else {
|
||
// 多个文件打包成zip
|
||
task.ImageLock.RUnlock()
|
||
var zipErr error
|
||
finalFilePath, zipErr = createZipArchive(task)
|
||
if zipErr != nil {
|
||
task.Status = StatusFailed
|
||
task.StatusLock.Unlock()
|
||
sendTaskUpdate(task)
|
||
return
|
||
}
|
||
}
|
||
|
||
if imageCount == 1 {
|
||
task.ImageLock.RUnlock()
|
||
}
|
||
|
||
task.OutputFile = finalFilePath
|
||
task.Status = StatusCompleted
|
||
|
||
// 设置完成计数为总任务数
|
||
task.ProgressLock.Lock()
|
||
task.CompletedCount = task.TotalCount
|
||
task.ProgressLock.Unlock()
|
||
|
||
task.StatusLock.Unlock()
|
||
|
||
// 发送最终状态更新
|
||
sendTaskUpdate(task)
|
||
|
||
// 确保所有进度都达到100%
|
||
ensureTaskCompletion(task)
|
||
|
||
fmt.Printf("任务 %s 全部完成: %d/%d\n", task.ID, task.CompletedCount, task.TotalCount)
|
||
}
|
||
|
||
// 下载单个镜像（带上下文控制）
|
||
func downloadImageWithContext(ctx context.Context, task *DownloadTask, index int, imgTask *ImageTask, platform string) error {
|
||
// 更新状态为运行中
|
||
sendProgressUpdate(task, index, 0, string(StatusRunning), "")
|
||
|
||
// 创建输出文件名
|
||
outputFileName := fmt.Sprintf("image_%d.tar", index)
|
||
outputPath := filepath.Join(task.TempDir, outputFileName)
|
||
|
||
imgTask.lock.Lock()
|
||
imgTask.OutputPath = outputPath
|
||
imgTask.lock.Unlock()
|
||
|
||
// 创建skopeo命令
|
||
platformArg := ""
|
||
if platform != "" {
|
||
// 支持手动输入完整的平台参数
|
||
if strings.Contains(platform, "--") {
|
||
platformArg = platform
|
||
} else {
|
||
// 处理特殊架构格式，如 arm/v7
|
||
if strings.Contains(platform, "/") {
|
||
parts := strings.Split(platform, "/")
|
||
if len(parts) == 2 {
|
||
// 适用于arm/v7这样的格式
|
||
platformArg = fmt.Sprintf("--override-os linux --override-arch %s --override-variant %s", parts[0], parts[1])
|
||
} else {
|
||
// 对于其他带/的格式，直接按原格式处理
|
||
platformArg = fmt.Sprintf("--override-os linux --override-arch %s", platform)
|
||
}
|
||
} else {
|
||
// 仅指定架构名称的情况
|
||
platformArg = fmt.Sprintf("--override-os linux --override-arch %s", platform)
|
||
}
|
||
}
|
||
}
|
||
|
||
// 构建命令
|
||
cmdStr := fmt.Sprintf("skopeo copy %s docker://%s docker-archive:%s",
|
||
platformArg, imgTask.Image, outputPath)
|
||
|
||
fmt.Printf("执行命令: %s\n", cmdStr)
|
||
|
||
// 创建可取消的命令
|
||
cmd := exec.CommandContext(ctx, "sh", "-c", cmdStr)
|
||
|
||
// 获取命令输出
|
||
stderr, err := cmd.StderrPipe()
|
||
if err != nil {
|
||
errMsg := fmt.Sprintf("无法创建输出管道: %v", err)
|
||
sendProgressUpdate(task, index, 0, string(StatusFailed), errMsg)
|
||
return fmt.Errorf(errMsg)
|
||
}
|
||
|
||
stdout, err := cmd.StdoutPipe()
|
||
if err != nil {
|
||
errMsg := fmt.Sprintf("无法创建标准输出管道: %v", err)
|
||
sendProgressUpdate(task, index, 0, string(StatusFailed), errMsg)
|
||
return fmt.Errorf(errMsg)
|
||
}
|
||
|
||
if err := cmd.Start(); err != nil {
|
||
errMsg := fmt.Sprintf("启动命令失败: %v", err)
|
||
sendProgressUpdate(task, index, 0, string(StatusFailed), errMsg)
|
||
return fmt.Errorf(errMsg)
|
||
}
|
||
|
||
// 使用进度通道传递进度信息
|
||
outputChan := make(chan string, 20)
|
||
done := make(chan struct{})
|
||
|
||
// 初始进度
|
||
sendProgressUpdate(task, index, 5, "", "")
|
||
|
||
// 进度聚合器
|
||
go func() {
|
||
// 镜像获取阶段的进度标记
|
||
downloadStages := map[string]float64{
|
||
"Getting image source signatures": 10,
|
||
"Copying blob": 30,
|
||
"Copying config": 70,
|
||
"Writing manifest": 90,
|
||
}
|
||
|
||
// 进度增长的定时器
|
||
ticker := time.NewTicker(100 * time.Millisecond)
|
||
defer ticker.Stop()
|
||
|
||
lastProgress := 5.0
|
||
stagnantTime := 0
|
||
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
// 上下文取消
|
||
return
|
||
|
||
case <-done:
|
||
// 命令完成，强制更新到100%
|
||
if lastProgress < 100 {
|
||
fmt.Printf("镜像 %s 下载完成，强制更新进度到100%%\n", imgTask.Image)
|
||
sendProgressUpdate(task, index, 100, string(StatusCompleted), "")
|
||
}
|
||
return
|
||
|
||
case output := <-outputChan:
|
||
// 解析输出更新进度
|
||
for marker, progress := range downloadStages {
|
||
if strings.Contains(output, marker) && progress > lastProgress {
|
||
lastProgress = progress
|
||
sendProgressUpdate(task, index, progress, "", "")
|
||
stagnantTime = 0
|
||
break
|
||
}
|
||
}
|
||
|
||
// 解析百分比
|
||
if strings.Contains(output, "%") {
|
||
parts := strings.Split(output, "%")
|
||
if len(parts) > 0 {
|
||
numStr := strings.TrimSpace(parts[0])
|
||
fields := strings.Fields(numStr)
|
||
if len(fields) > 0 {
|
||
lastField := fields[len(fields)-1]
|
||
parsedProgress := 0.0
|
||
_, err := fmt.Sscanf(lastField, "%f", &parsedProgress)
|
||
if err == nil && parsedProgress > 0 && parsedProgress <= 100 {
|
||
// 根据当前阶段调整进度值
|
||
var adjustedProgress float64
|
||
if lastProgress < 30 {
|
||
// Copying blob阶段，进度在10-30%之间
|
||
adjustedProgress = 10 + (parsedProgress / 100) * 20
|
||
} else if lastProgress < 70 {
|
||
// Copying config阶段，进度在30-70%之间
|
||
adjustedProgress = 30 + (parsedProgress / 100) * 40
|
||
} else if lastProgress < 90 {
|
||
// Writing manifest阶段，进度在70-90%之间
|
||
adjustedProgress = 70 + (parsedProgress / 100) * 20
|
||
}
|
||
|
||
if adjustedProgress > lastProgress {
|
||
lastProgress = adjustedProgress
|
||
sendProgressUpdate(task, index, adjustedProgress, "", "")
|
||
stagnantTime = 0
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// 如果发现完成标记，立即更新到100%
|
||
if checkForCompletionMarkers(output) {
|
||
fmt.Printf("镜像 %s 检测到完成标记\n", imgTask.Image)
|
||
lastProgress = 100
|
||
sendProgressUpdate(task, index, 100, string(StatusCompleted), "")
|
||
stagnantTime = 0
|
||
}
|
||
|
||
case <-ticker.C:
|
||
// 如果进度长时间无变化，缓慢增加
|
||
stagnantTime += 100 // 100ms
|
||
if stagnantTime >= 10000 && lastProgress < 95 { // 10秒无变化
|
||
// 每10秒增加5%进度，确保不超过95%
|
||
newProgress := lastProgress + 5
|
||
if newProgress > 95 {
|
||
newProgress = 95
|
||
}
|
||
lastProgress = newProgress
|
||
sendProgressUpdate(task, index, newProgress, "", "")
|
||
stagnantTime = 0
|
||
}
|
||
}
|
||
}
|
||
}()
|
||
|
||
// 读取标准输出
|
||
go func() {
|
||
scanner := bufio.NewScanner(stdout)
|
||
for scanner.Scan() {
|
||
output := scanner.Text()
|
||
fmt.Printf("镜像 %s 标准输出: %s\n", imgTask.Image, output)
|
||
select {
|
||
case outputChan <- output:
|
||
default:
|
||
// 通道已满，丢弃
|
||
}
|
||
}
|
||
}()
|
||
|
||
// 读取错误输出
|
||
go func() {
|
||
scanner := bufio.NewScanner(stderr)
|
||
for scanner.Scan() {
|
||
output := scanner.Text()
|
||
fmt.Printf("镜像 %s 错误输出: %s\n", imgTask.Image, output)
|
||
select {
|
||
case outputChan <- output:
|
||
default:
|
||
// 通道已满，丢弃
|
||
}
|
||
}
|
||
}()
|
||
|
||
// 等待命令完成
|
||
cmdErr := cmd.Wait()
|
||
close(done) // 通知进度聚合器退出
|
||
|
||
if cmdErr != nil {
|
||
errMsg := fmt.Sprintf("命令执行失败: %v", cmdErr)
|
||
sendProgressUpdate(task, index, 0, string(StatusFailed), errMsg)
|
||
return fmt.Errorf(errMsg)
|
||
}
|
||
|
||
// 检查文件是否成功创建
|
||
if _, err := os.Stat(outputPath); os.IsNotExist(err) {
|
||
errMsg := "文件未成功创建"
|
||
sendProgressUpdate(task, index, 0, string(StatusFailed), errMsg)
|
||
return fmt.Errorf(errMsg)
|
||
}
|
||
|
||
// 确保更新状态为已完成，进度为100%
|
||
sendProgressUpdate(task, index, 100, string(StatusCompleted), "")
|
||
return nil
|
||
}
|
||
|
||
// 创建ZIP归档
|
||
func createZipArchive(task *DownloadTask) (string, error) {
|
||
zipFilePath := filepath.Join(task.TempDir, "images.zip")
|
||
zipFile, err := os.Create(zipFilePath)
|
||
if err != nil {
|
||
return "", fmt.Errorf("创建ZIP文件失败: %w", err)
|
||
}
|
||
defer zipFile.Close()
|
||
|
||
zipWriter := zip.NewWriter(zipFile)
|
||
defer zipWriter.Close()
|
||
|
||
task.ImageLock.RLock()
|
||
images := make([]*ImageTask, len(task.Images))
|
||
copy(images, task.Images) // 创建副本避免长时间持有锁
|
||
task.ImageLock.RUnlock()
|
||
|
||
for _, img := range images {
|
||
img.lock.Lock()
|
||
status := img.Status
|
||
outputPath := img.OutputPath
|
||
image := img.Image
|
||
img.lock.Unlock()
|
||
|
||
if status != string(StatusCompleted) || outputPath == "" {
|
||
continue
|
||
}
|
||
|
||
// 创建ZIP条目
|
||
imgFile, err := os.Open(outputPath)
|
||
if err != nil {
|
||
return "", fmt.Errorf("无法打开镜像文件 %s: %w", outputPath, err)
|
||
}
|
||
|
||
// 使用镜像名作为文件名
|
||
imageName := strings.ReplaceAll(image, "/", "_")
|
||
imageName = strings.ReplaceAll(imageName, ":", "_")
|
||
fileName := imageName + ".tar"
|
||
|
||
fileInfo, err := imgFile.Stat()
|
||
if err != nil {
|
||
imgFile.Close()
|
||
return "", fmt.Errorf("无法获取文件信息: %w", err)
|
||
}
|
||
|
||
header, err := zip.FileInfoHeader(fileInfo)
|
||
if err != nil {
|
||
imgFile.Close()
|
||
return "", fmt.Errorf("创建ZIP头信息失败: %w", err)
|
||
}
|
||
|
||
header.Name = fileName
|
||
header.Method = zip.Deflate
|
||
|
||
writer, err := zipWriter.CreateHeader(header)
|
||
if err != nil {
|
||
imgFile.Close()
|
||
return "", fmt.Errorf("添加文件到ZIP失败: %w", err)
|
||
}
|
||
|
||
_, err = io.Copy(writer, imgFile)
|
||
imgFile.Close()
|
||
if err != nil {
|
||
return "", fmt.Errorf("写入ZIP文件失败: %w", err)
|
||
}
|
||
}
|
||
|
||
return zipFilePath, nil
|
||
}
|
||
|
||
// 发送任务更新到WebSocket
|
||
func sendTaskUpdate(task *DownloadTask) {
|
||
// 复制任务状态避免序列化时锁定
|
||
taskCopy := &DownloadTask{
|
||
ID: task.ID,
|
||
CompletedCount: 0,
|
||
TotalCount: len(task.Images),
|
||
Status: TaskStatus(""),
|
||
Images: nil,
|
||
}
|
||
|
||
// 复制状态信息
|
||
task.StatusLock.RLock()
|
||
taskCopy.Status = task.Status
|
||
task.StatusLock.RUnlock()
|
||
|
||
task.ProgressLock.RLock()
|
||
taskCopy.CompletedCount = task.CompletedCount
|
||
task.ProgressLock.RUnlock()
|
||
|
||
// 复制镜像信息
|
||
task.ImageLock.RLock()
|
||
taskCopy.Images = make([]*ImageTask, len(task.Images))
|
||
for i, img := range task.Images {
|
||
img.lock.Lock()
|
||
taskCopy.Images[i] = &ImageTask{
|
||
Image: img.Image,
|
||
Progress: img.Progress,
|
||
Status: img.Status,
|
||
Error: img.Error,
|
||
}
|
||
img.lock.Unlock()
|
||
}
|
||
task.ImageLock.RUnlock()
|
||
|
||
// 序列化并发送
|
||
taskJSON, err := json.Marshal(taskCopy)
|
||
if err != nil {
|
||
fmt.Printf("序列化任务失败: %v\n", err)
|
||
return
|
||
}
|
||
|
||
clientLock.Lock()
|
||
client, exists := clients[task.ID]
|
||
clientLock.Unlock()
|
||
|
||
if exists {
|
||
select {
|
||
case client.Send <- taskJSON:
|
||
// 成功发送
|
||
default:
|
||
// 通道已满或关闭，忽略
|
||
}
|
||
}
|
||
}
|
||
|
||
// 发送单个镜像更新 - 保持兼容性
|
||
func sendImageUpdate(task *DownloadTask, imageIndex int) {
|
||
sendTaskUpdate(task)
|
||
}
|
||
|
||
// 提供文件下载
|
||
func serveFile(c *gin.Context) {
|
||
filename := c.Param("filename")
|
||
|
||
// 安全检查，防止任意文件访问
|
||
if strings.Contains(filename, "..") {
|
||
c.JSON(http.StatusForbidden, gin.H{"error": "无效的文件名"})
|
||
return
|
||
}
|
||
|
||
// 根据任务ID和文件名查找文件
|
||
parts := strings.Split(filename, "_")
|
||
if len(parts) < 2 {
|
||
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的文件名格式"})
|
||
return
|
||
}
|
||
|
||
taskID := parts[0]
|
||
|
||
tasksLock.Lock()
|
||
task, exists := tasks[taskID]
|
||
tasksLock.Unlock()
|
||
|
||
if !exists {
|
||
c.JSON(http.StatusNotFound, gin.H{"error": "任务不存在"})
|
||
return
|
||
}
|
||
|
||
// 确保任务状态为已完成，并且所有进度都是100%
|
||
task.StatusLock.RLock()
|
||
isCompleted := task.Status == StatusCompleted
|
||
task.StatusLock.RUnlock()
|
||
|
||
if isCompleted {
|
||
// 确保所有进度达到100%
|
||
ensureTaskCompletion(task)
|
||
}
|
||
|
||
// 检查文件是否存在
|
||
filePath := task.OutputFile
|
||
if filePath == "" || !fileExists(filePath) {
|
||
c.JSON(http.StatusNotFound, gin.H{"error": "文件不存在"})
|
||
return
|
||
}
|
||
|
||
// 获取文件信息
|
||
fileInfo, err := os.Stat(filePath)
|
||
if err != nil {
|
||
c.JSON(http.StatusInternalServerError, gin.H{"error": "无法获取文件信息"})
|
||
return
|
||
}
|
||
|
||
// 设置文件名 - 提取有意义的文件名
|
||
downloadName := filepath.Base(filePath)
|
||
c.Header("Content-Disposition", fmt.Sprintf("attachment; filename=%s", downloadName))
|
||
c.Header("Content-Length", fmt.Sprintf("%d", fileInfo.Size()))
|
||
|
||
// 返回文件
|
||
c.File(filePath)
|
||
}
|
||
|
||
// 检查文件是否存在
|
||
func fileExists(path string) bool {
|
||
_, err := os.Stat(path)
|
||
return !os.IsNotExist(err)
|
||
}
|
||
|
||
// 清理过期临时文件
|
||
func cleanupTempFiles() {
|
||
// 创建两个定时器
|
||
hourlyTicker := time.NewTicker(1 * time.Hour)
|
||
fiveMinTicker := time.NewTicker(5 * time.Minute)
|
||
|
||
// 清理所有文件的函数
|
||
cleanAll := func() {
|
||
fmt.Printf("执行清理所有临时文件\n")
|
||
entries, err := os.ReadDir("./temp")
|
||
if err == nil {
|
||
for _, entry := range entries {
|
||
entryPath := filepath.Join("./temp", entry.Name())
|
||
info, err := entry.Info()
|
||
if err == nil {
|
||
if info.IsDir() {
|
||
os.RemoveAll(entryPath)
|
||
} else {
|
||
os.Remove(entryPath)
|
||
}
|
||
}
|
||
}
|
||
} else {
|
||
fmt.Printf("清理临时文件失败: %v\n", err)
|
||
}
|
||
}
|
||
|
||
// 检查文件大小并在需要时清理
|
||
checkSizeAndClean := func() {
|
||
var totalSize int64 = 0
|
||
err := filepath.Walk("./temp", func(path string, info os.FileInfo, err error) error {
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
// 跳过根目录
|
||
if path == "./temp" {
|
||
return nil
|
||
}
|
||
|
||
if !info.IsDir() {
|
||
totalSize += info.Size()
|
||
}
|
||
|
||
return nil
|
||
})
|
||
|
||
if err != nil {
|
||
fmt.Printf("计算临时文件总大小失败: %v\n", err)
|
||
return
|
||
}
|
||
|
||
// 如果总大小超过10GB，清理所有文件，防止恶意下载导致磁盘爆满
|
||
if totalSize > 10*1024*1024*1024 {
|
||
fmt.Printf("临时文件总大小超过10GB (当前: %.2f GB)，清理所有文件\n", float64(totalSize)/(1024*1024*1024))
|
||
cleanAll()
|
||
} else {
|
||
fmt.Printf("临时文件总大小: %.2f GB\n", float64(totalSize)/(1024*1024*1024))
|
||
}
|
||
}
|
||
|
||
// 主循环
|
||
for {
|
||
select {
|
||
case <-hourlyTicker.C:
|
||
// 每小时清理所有文件
|
||
cleanAll()
|
||
case <-fiveMinTicker.C:
|
||
// 每5分钟检查一次总文件大小
|
||
checkSizeAndClean()
|
||
}
|
||
}
|
||
}
|
||
|
||
// 完成任务处理函数，确保进度是100%
|
||
func ensureTaskCompletion(task *DownloadTask) {
|
||
// 重新检查一遍所有镜像的进度
|
||
task.ImageLock.RLock()
|
||
completedCount := 0
|
||
totalCount := len(task.Images)
|
||
|
||
for i, img := range task.Images {
|
||
img.lock.Lock()
|
||
if img.Status == string(StatusCompleted) {
|
||
// 确保进度为100%
|
||
if img.Progress < 100 {
|
||
img.Progress = 100
|
||
fmt.Printf("确保镜像 %d 进度为100%%\n", i)
|
||
}
|
||
completedCount++
|
||
}
|
||
img.lock.Unlock()
|
||
}
|
||
task.ImageLock.RUnlock()
|
||
|
||
// 更新完成计数
|
||
task.ProgressLock.Lock()
|
||
task.CompletedCount = completedCount
|
||
task.TotalCount = totalCount
|
||
task.ProgressLock.Unlock()
|
||
|
||
// 如果任务状态为已完成，但计数不匹配，修正计数
|
||
task.StatusLock.RLock()
|
||
isCompleted := task.Status == StatusCompleted
|
||
task.StatusLock.RUnlock()
|
||
|
||
if isCompleted && completedCount != totalCount {
|
||
task.ProgressLock.Lock()
|
||
task.CompletedCount = totalCount
|
||
task.ProgressLock.Unlock()
|
||
fmt.Printf("任务 %s 状态已完成，强制设置计数为 %d/%d\n", task.ID, totalCount, totalCount)
|
||
}
|
||
|
||
// 发送最终更新
|
||
sendTaskUpdate(task)
|
||
}
|
||
|
||
// 处理下载单个镜像的输出中的完成标记
|
||
func checkForCompletionMarkers(output string) bool {
|
||
// 已知的完成标记
|
||
completionMarkers := []string{
|
||
"Writing manifest to image destination",
|
||
"Copying config complete",
|
||
"Storing signatures",
|
||
"Writing manifest complete",
|
||
}
|
||
|
||
for _, marker := range completionMarkers {
|
||
if strings.Contains(output, marker) {
|
||
return true
|
||
}
|
||
}
|
||
|
||
return false
|
||
} |