cmsp/internal/downloader/downloader.go

// Package downloader 提供单个视频的下载和 ffprobe 完整性校验。
package downloader

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"
	"unicode"
)

const downloadTimeout = 180 * time.Second

type ProbeResult struct {
	Duration   float64
	Size       int64
	FormatName string
}

type ProbeFunc func(context.Context, string) (ProbeResult, error)

type Result struct {
	Path     string
	Size     int64
	Duration float64
	Skipped  bool
}

type Downloader struct {
	client *http.Client
	probe  ProbeFunc
}

func New() *Downloader {
	return &Downloader{
		client: &http.Client{Timeout: downloadTimeout},
		probe:  runFFProbe,
	}
}

// NewWithOptions 只用于测试替换 HTTP 客户端和 ffprobe 执行。
func NewWithOptions(client *http.Client, probe ProbeFunc) *Downloader {
	if client == nil {
		client = &http.Client{Timeout: downloadTimeout}
	}
	if probe == nil {
		probe = runFFProbe
	}
	return &Downloader{client: client, probe: probe}
}

// Filename 生成 Windows 可用且不会逃出目标目录的视频文件名。
//
// 形如 40583431295_1.mp4：前缀是蝦皮商品 ID（和所在子目录同名），
// 后面是该商品的视频序号，从 1 开始。
//
// 文件被单独挪走或和别的商品混在一起时，靠前缀仍能认出属于哪个商品，
// 后续别的程序也好按前缀匹配。
//
// 视频来自哪个淘宝同款不体现在文件名里，那个记录在 videos.source_item，
// 需要回溯来源时查库。
func Filename(shopeeItemID string, index int) string {
	if index < 1 {
		index = 1
	}
	return fmt.Sprintf("%s_%d.mp4", safeID(shopeeItemID), index)
}

func safeID(value string) string {
	var b strings.Builder
	for _, r := range strings.TrimSpace(value) {
		if unicode.IsLetter(r) || unicode.IsDigit(r) || r == '-' || r == '_' {
			b.WriteRune(r)
		} else {
			b.WriteByte('_')
		}
	}
	value = strings.Trim(b.String(), " .")
	if value == "" {
		return "unknown"
	}
	// 全是占位下划线说明原值里没有任何可用字符（比如 ".." 或 "///"），
	// 这种目录名毫无意义，回落到 unknown 更好排查。
	if strings.Trim(value, "_") == "" {
		return "unknown"
	}
	return value
}

// Download 下载到 .part，ffprobe 校验成功后才改名为正式文件。
func (d *Downloader) Download(ctx context.Context, sourceURL, referer, target string, retries int) (Result, error) {
	if info, err := os.Stat(target); err == nil && !info.IsDir() && info.Size() > 0 {
		return Result{Path: target, Size: info.Size(), Skipped: true}, nil
	}
	if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
		return Result{}, fmt.Errorf("创建视频目录失败：%w", err)
	}
	part := target + ".part"
	if retries < 0 {
		retries = 0
	}
	// 超时按「每次尝试」计，和参考实现的 curl --max-time 180 一致：
	// curl 每次重试都是一次独立调用，各有完整 180 秒。
	// 若把超时套在整个重试循环外面，大视频跑到一半失败后剩余预算不足，
	// 重试必然再次超时，等于没有重试。
	for attempt := 0; attempt <= retries; attempt++ {
		result, retryable, err := d.attemptDownload(ctx, sourceURL, referer, target, part)
		if err == nil {
			return result, nil
		}
		if !retryable || attempt == retries {
			return Result{}, err
		}
		wait := time.Second << attempt
		timer := time.NewTimer(wait)
		select {
		case <-ctx.Done():
			timer.Stop()
			return Result{}, ctx.Err()
		case <-timer.C:
		}
	}
	return Result{}, fmt.Errorf("下载视频请求失败")
}

// attemptDownload 为单次尝试套上独立的 180 秒超时。父 context 被取消时
// 仍然立即中止，停止任务不会卡在这里。
func (d *Downloader) attemptDownload(ctx context.Context, sourceURL, referer, target, part string) (Result, bool, error) {
	attemptCtx, cancel := context.WithTimeout(ctx, downloadTimeout)
	defer cancel()
	return d.downloadOnce(attemptCtx, sourceURL, referer, target, part)
}

// downloadOnce 只把可由网络恢复的失败标为 retryable；内容校验失败绝不重试。
func (d *Downloader) downloadOnce(ctx context.Context, sourceURL, referer, target, part string) (Result, bool, error) {
	_ = os.Remove(part)

	request, err := http.NewRequestWithContext(ctx, http.MethodGet, sourceURL, nil)
	if err != nil {
		return Result{}, false, fmt.Errorf("视频地址格式无效")
	}
	request.Header.Set("User-Agent", "Mozilla/5.0")
	request.Header.Set("Referer", referer)
	response, err := d.client.Do(request)
	if err != nil {
		// net/http 的错误通常包含完整 URL（包括查询签名），不能向日志上抛。
		return Result{}, true, fmt.Errorf("下载视频请求失败")
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		_, _ = io.Copy(io.Discard, response.Body)
		return Result{}, response.StatusCode >= 500, fmt.Errorf("下载视频失败：HTTP %d", response.StatusCode)
	}

	file, err := os.OpenFile(part, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
	if err != nil {
		return Result{}, false, fmt.Errorf("创建视频临时文件失败：%w", err)
	}
	_, copyErr := io.Copy(file, response.Body)
	closeErr := file.Close()
	if copyErr != nil {
		_ = os.Remove(part)
		return Result{}, true, fmt.Errorf("写入视频临时文件失败：%w", copyErr)
	}
	if closeErr != nil {
		_ = os.Remove(part)
		return Result{}, false, fmt.Errorf("关闭视频临时文件失败：%w", closeErr)
	}

	probe, err := d.probe(ctx, part)
	if err != nil {
		_ = os.Remove(part)
		return Result{}, false, fmt.Errorf("视频完整性校验失败：%w", err)
	}
	info, err := os.Stat(part)
	if err != nil {
		_ = os.Remove(part)
		return Result{}, false, fmt.Errorf("读取视频临时文件失败：%w", err)
	}
	size := probe.Size
	if size <= 0 {
		size = info.Size()
	}
	if size <= 0 {
		_ = os.Remove(part)
		return Result{}, false, fmt.Errorf("视频完整性校验失败：文件大小为 0")
	}
	if err := os.Remove(target); err != nil && !os.IsNotExist(err) {
		_ = os.Remove(part)
		return Result{}, false, fmt.Errorf("替换空目标文件失败：%w", err)
	}
	if err := os.Rename(part, target); err != nil {
		_ = os.Remove(part)
		return Result{}, false, fmt.Errorf("保存正式视频文件失败：%w", err)
	}
	return Result{Path: target, Size: size, Duration: probe.Duration}, false, nil
}

func runFFProbe(ctx context.Context, path string) (ProbeResult, error) {
	ffprobe, err := exec.LookPath("ffprobe")
	if err != nil {
		return ProbeResult{}, fmt.Errorf("未找到 ffprobe，请先安装并加入 PATH")
	}
	command := exec.CommandContext(ctx, ffprobe,
		"-v", "error", "-show_entries", "format=duration,size,format_name", "-of", "json", path)
	output, err := command.Output()
	if err != nil {
		return ProbeResult{}, fmt.Errorf("ffprobe 执行失败：%w", err)
	}
	var payload struct {
		Format struct {
			Duration   string `json:"duration"`
			Size       string `json:"size"`
			FormatName string `json:"format_name"`
		} `json:"format"`
	}
	if err := json.Unmarshal(output, &payload); err != nil {
		return ProbeResult{}, fmt.Errorf("解析 ffprobe 输出失败：%w", err)
	}
	duration, err := strconv.ParseFloat(payload.Format.Duration, 64)
	if err != nil || duration <= 0 {
		return ProbeResult{}, fmt.Errorf("ffprobe 未返回有效 duration")
	}
	size, err := strconv.ParseInt(payload.Format.Size, 10, 64)
	if err != nil || size <= 0 {
		return ProbeResult{}, fmt.Errorf("ffprobe 未返回有效 size")
	}
	if strings.TrimSpace(payload.Format.FormatName) == "" {
		return ProbeResult{}, fmt.Errorf("ffprobe 未返回 format_name")
	}
	return ProbeResult{Duration: duration, Size: size, FormatName: payload.Format.FormatName}, nil
}

// SafeDirName 把蝦皮商品 ID 清洗成可以直接当目录名的字符串。
//
// 商品 ID 正常是纯数字，但不能假定：一旦货憨憨返回带路径分隔符或
// .. 的值，直接拼进路径就能写到目录之外。这里统一只保留字母数字
// 和减号下划线。
func SafeDirName(value string) string {
	return safeID(value)
}