mirror of
https://github.com/sky22333/hubproxy.git
synced 2026-08-09 21:52:11 +08:00
1114 lines
27 KiB
Go
1114 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
|
|
}
|
|
|
|
// 创建新任务
|
|
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 { // 15GB
|
|
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
|
|
} |