update: chunk upload

2026-09-11 20:02:47 +08:00
parent 51f1fe96bf
commit d225b6e182
7 changed files with 584 additions and 13 deletions
@@ -0,0 +1,366 @@
+package main
+
+import (
+	"encoding/json"
+	"io"
+	"log"
+	"net/http"
+	"os"
+	"path/filepath"
+	"regexp"
+	"strconv"
+	"sync"
+	"time"
+
+	"github.com/gin-gonic/gin"
+)
+
+const (
+	chunkDir       = "./data/chunks"
+	chunkThreshold = 5 << 30 // >5GB 才走分片
+	chunkSize      = 1 << 30 // 每片 1GB
+)
+
+var uploadIDRe = regexp.MustCompile(`^[a-zA-Z0-9_-]{8,64}$`)
+
+type chunkMeta struct {
+	UploadID    string       `json:"upload_id"`
+	Filename    string       `json:"filename"`
+	TotalSize   int64        `json:"total_size"`
+	TotalChunks int          `json:"total_chunks"`
+	Received    map[int]bool `json:"received"`
+	CreatedAt   time.Time    `json:"created_at"`
+	Reserved    int64        `json:"reserved"`
+}
+
+// 全局 map 存 in-flight 上传; 持锁只做元数据操作, 1GB 分片 io.Copy 在锁外进行.
+// ponytail: 全局锁, 若并发大文件上传成为瓶颈再换 per-upload 锁.
+var (
+	chunkMu      sync.RWMutex
+	chunkMetaMap = map[string]*chunkMeta{}
+)
+
+func chunkDirPath(id string) string { return filepath.Join(chunkDir, id) }
+func chunkFilePath(id string, idx int) string {
+	return filepath.Join(chunkDir, id, strconv.Itoa(idx))
+}
+
+func metaPath(id string) string { return filepath.Join(chunkDir, id, "meta.json") }
+
+// saveMeta 原子写 meta.json (.tmp + rename).
+func saveMeta(m *chunkMeta) error {
+	b, err := json.Marshal(m)
+	if err != nil {
+		return err
+	}
+	tmp := metaPath(m.UploadID) + ".tmp"
+	if err := os.WriteFile(tmp, b, 0o644); err != nil {
+		return err
+	}
+	return os.Rename(tmp, metaPath(m.UploadID))
+}
+
+func loadMeta(id string) *chunkMeta {
+	b, err := os.ReadFile(metaPath(id))
+	if err != nil {
+		return nil
+	}
+	var m chunkMeta
+	if err := json.Unmarshal(b, &m); err != nil {
+		return nil
+	}
+	return &m
+}
+
+// expectedChunks totalSize 按 chunkSize 分片应有的片数
+func expectedChunks(totalSize int64) int {
+	return int((totalSize + chunkSize - 1) / chunkSize)
+}
+
+// expectedChunkSize 第 i 片应有的字节数 (最后一片是余数)
+func expectedChunkSize(totalSize int64, i int) int64 {
+	if n := totalSize - int64(i)*chunkSize; n < chunkSize {
+		return n
+	}
+	return chunkSize
+}
+
+func uploadChunkHandler(c *gin.Context) {
+	uploadID := c.GetHeader("X-Upload-Id")
+	if !uploadIDRe.MatchString(uploadID) {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid upload id"})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "invalid_upload_id"})
+		return
+	}
+	idx, err := strconv.Atoi(c.GetHeader("X-Chunk-Index"))
+	if err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid chunk index"})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "invalid_chunk_index"})
+		return
+	}
+	totalChunks, err := strconv.Atoi(c.GetHeader("X-Total-Chunks"))
+	if err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid total chunks"})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "invalid_total_chunks"})
+		return
+	}
+	filename := c.GetHeader("X-Filename")
+	if !checkFilenameSafe(filename) {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid filename"})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "invalid_filename"})
+		return
+	}
+	totalSize, err := strconv.ParseInt(c.GetHeader("X-Total-Size"), 10, 64)
+	if err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid total size"})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "invalid_total_size"})
+		return
+	}
+	if totalSize <= chunkThreshold {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "chunked upload is for files > 5GB"})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "below_threshold"})
+		return
+	}
+	if totalChunks < 1 || idx < 0 || idx >= totalChunks {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "chunk index out of range"})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "index": idx, "total": totalChunks, "result": "index_out_of_range"})
+		return
+	}
+	if e := expectedChunks(totalSize); totalChunks < e-1 || totalChunks > e+1 {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "total chunks mismatch"})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "total": totalChunks, "expected": e, "result": "chunk_count_mismatch"})
+		return
+	}
+	if totalSize > quotaRemaining() {
+		c.JSON(http.StatusInsufficientStorage, gin.H{"error": "quota exceeded"})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "size": totalSize, "result": "quota_exceeded"})
+		return
+	}
+
+	chunk, err := c.FormFile("chunk")
+	if err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "missing form field 'chunk'"})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "missing_form_field"})
+		return
+	}
+
+	dir := chunkDirPath(uploadID)
+	chunkMu.RLock()
+	meta, inMap := chunkMetaMap[uploadID]
+	chunkMu.RUnlock()
+	if !inMap {
+		meta = loadMeta(uploadID) // 进程重启后从磁盘恢复
+		if meta != nil {
+			chunkMu.Lock()
+			chunkMetaMap[uploadID] = meta
+			chunkMu.Unlock()
+		}
+	}
+	if meta == nil {
+		if idx != 0 {
+			c.JSON(http.StatusNotFound, gin.H{"error": "upload not found"})
+			audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "upload_not_found"})
+			return
+		}
+		if err := os.MkdirAll(dir, 0o755); err != nil {
+			c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
+			audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "mkdir_failed"})
+			return
+		}
+		if !quotaReserve(totalSize) {
+			_ = os.RemoveAll(dir)
+			c.JSON(http.StatusInsufficientStorage, gin.H{"error": "quota exceeded"})
+			audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "size": totalSize, "result": "quota_exceeded"})
+			return
+		}
+		meta = &chunkMeta{
+			UploadID: uploadID, Filename: filename, TotalSize: totalSize,
+			TotalChunks: totalChunks, Received: map[int]bool{},
+			CreatedAt: time.Now(), Reserved: totalSize,
+		}
+		chunkMu.Lock()
+		chunkMetaMap[uploadID] = meta
+		err = saveMeta(meta)
+		chunkMu.Unlock()
+		if err != nil {
+			abortChunkUpload(uploadID, meta)
+			c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
+			audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "meta_save_failed"})
+			return
+		}
+	} else if idx > 0 {
+		if _, err := os.Stat(dir); err != nil {
+			c.JSON(http.StatusNotFound, gin.H{"error": "upload not found"})
+			audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "result": "upload_not_found"})
+			return
+		}
+	}
+
+	// 写分片: multipart → 流式 io.Copy 落盘, 不进内存
+	dst := chunkFilePath(uploadID, idx)
+	in, err := chunk.Open()
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "index": idx, "result": "open_failed"})
+		return
+	}
+	out, err := os.Create(dst)
+	if err != nil {
+		in.Close()
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "index": idx, "result": "create_failed"})
+		return
+	}
+	_, cpErr := io.Copy(out, in)
+	in.Close()
+	if cerr := out.Close(); cerr != nil && cpErr == nil {
+		cpErr = cerr
+	}
+	if cpErr != nil {
+		_ = os.Remove(dst)
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "index": idx, "result": "write_failed"})
+		return
+	}
+
+	// 更新 Received + 持久化 (持写锁, 防并发分片时 meta.json 丢进度)
+	chunkMu.Lock()
+	meta.Received[idx] = true
+	err = saveMeta(meta)
+	complete := len(meta.Received) == meta.TotalChunks
+	chunkMu.Unlock()
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "index": idx, "result": "meta_save_failed"})
+		return
+	}
+
+	if !complete {
+		c.JSON(http.StatusAccepted, gin.H{"received": len(meta.Received), "total": meta.TotalChunks})
+		audit(c, "upload_chunk", gin.H{"upload_id": uploadID, "index": idx, "received": len(meta.Received), "total": meta.TotalChunks, "result": "ok"})
+		return
+	}
+	assembleChunks(c, uploadID, meta)
+}
+
+// assembleChunks 所有分片到位后顺序拼装落盘并清理.
+// ponytail: 末片重传/并发组装不额外去重, 前端单分片串行重试不会触发.
+func assembleChunks(c *gin.Context, uploadID string, meta *chunkMeta) {
+	dir := chunkDirPath(uploadID)
+	for i := 0; i < meta.TotalChunks; i++ {
+		info, err := os.Stat(chunkFilePath(uploadID, i))
+		if err != nil || info.Size() != expectedChunkSize(meta.TotalSize, i) {
+			abortChunkUpload(uploadID, meta)
+			c.JSON(http.StatusBadRequest, gin.H{"error": "chunk size mismatch"})
+			audit(c, "upload_assembled", gin.H{"upload_id": uploadID, "index": i, "result": "chunk_size_mismatch"})
+			return
+		}
+	}
+
+	name := uniqueFilename(*uploadDir, meta.Filename)
+	dst := filepath.Join(*uploadDir, name)
+	out, err := os.Create(dst)
+	if err != nil {
+		abortChunkUpload(uploadID, meta)
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
+		audit(c, "upload_assembled", gin.H{"upload_id": uploadID, "result": "create_failed"})
+		return
+	}
+	var realSize int64
+	for i := 0; i < meta.TotalChunks; i++ {
+		in, err := os.Open(chunkFilePath(uploadID, i))
+		if err != nil {
+			out.Close()
+			_ = os.Remove(dst)
+			abortChunkUpload(uploadID, meta)
+			c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
+			audit(c, "upload_assembled", gin.H{"upload_id": uploadID, "index": i, "result": "open_failed"})
+			return
+		}
+		n, cErr := io.Copy(out, in)
+		in.Close()
+		if cErr != nil {
+			out.Close()
+			_ = os.Remove(dst)
+			abortChunkUpload(uploadID, meta)
+			c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
+			audit(c, "upload_assembled", gin.H{"upload_id": uploadID, "index": i, "result": "copy_failed"})
+			return
+		}
+		realSize += n
+	}
+	if err := out.Close(); err != nil {
+		_ = os.Remove(dst)
+		abortChunkUpload(uploadID, meta)
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error"})
+		audit(c, "upload_assembled", gin.H{"upload_id": uploadID, "result": "close_failed"})
+		return
+	}
+
+	// 配额: 释放预留, 按真实大小重新占用
+	quotaSub(meta.Reserved)
+	quotaReserve(realSize)
+	chunkMu.Lock()
+	delete(chunkMetaMap, uploadID)
+	chunkMu.Unlock()
+	_ = os.RemoveAll(dir)
+
+	now := time.Now()
+	_ = os.Chtimes(dst, now, now)
+	c.JSON(http.StatusOK, fileResponse{
+		Filename: name, URL: "/download/" + name, Size: realSize,
+		UploadedAt: now, ExpiresAt: now.Add(*fileTTL),
+	})
+	audit(c, "upload_assembled", gin.H{
+		"upload_id": uploadID, "file": name, "size": realSize, "result": "ok",
+	})
+}
+
+// abortChunkUpload 失败时释放预留配额并清掉分片目录与内存记录
+func abortChunkUpload(uploadID string, meta *chunkMeta) {
+	quotaSub(meta.Reserved)
+	chunkMu.Lock()
+	delete(chunkMetaMap, uploadID)
+	chunkMu.Unlock()
+	_ = os.RemoveAll(chunkDirPath(uploadID))
+}
+
+// cleanupChunks 清理超过 maxAge 的孤儿分片目录并释放预留配额
+func cleanupChunks(maxAge time.Duration) {
+	entries, err := os.ReadDir(chunkDir)
+	if err != nil {
+		if !os.IsNotExist(err) {
+			log.Printf("[cleaner] read chunk dir failed: %v", err)
+		}
+		return
+	}
+	now := time.Now()
+	for _, e := range entries {
+		if !e.IsDir() {
+			continue
+		}
+		id := e.Name()
+		dir := chunkDirPath(id)
+		if !uploadIDRe.MatchString(id) {
+			_ = os.RemoveAll(dir) // 非法目录名 (非本服务生成), 直接清
+			continue
+		}
+		meta := loadMeta(id)
+		if meta == nil {
+			// 无 meta 的孤儿目录: 无法确定 reserved, 删目录即可 (quota 由 cleaner 的 rescanQuota 校正)
+			_ = os.RemoveAll(dir)
+			chunkMu.Lock()
+			delete(chunkMetaMap, id)
+			chunkMu.Unlock()
+			continue
+		}
+		if now.Sub(meta.CreatedAt) > maxAge {
+			quotaSub(meta.Reserved)
+			chunkMu.Lock()
+			delete(chunkMetaMap, id)
+			chunkMu.Unlock()
+			_ = os.RemoveAll(dir)
+			log.Printf("[cleaner] removed orphan chunks: %s (age=%v)", id, now.Sub(meta.CreatedAt))
+		}
+	}
+}