github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/thumb/ffmpeg.go (about)

     1  package thumb
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"fmt"
     7  	model "github.com/cloudreve/Cloudreve/v3/models"
     8  	"github.com/cloudreve/Cloudreve/v3/pkg/util"
     9  	"github.com/gofrs/uuid"
    10  	"io"
    11  	"os"
    12  	"os/exec"
    13  	"path/filepath"
    14  	"strings"
    15  )
    16  
    17  func init() {
    18  	RegisterGenerator(&FfmpegGenerator{})
    19  }
    20  
    21  type FfmpegGenerator struct {
    22  	exts        []string
    23  	lastRawExts string
    24  }
    25  
    26  func (f *FfmpegGenerator) Generate(ctx context.Context, file io.Reader, src, name string, options map[string]string) (*Result, error) {
    27  	ffmpegOpts := model.GetSettingByNames("thumb_ffmpeg_path", "thumb_ffmpeg_exts", "thumb_ffmpeg_seek", "thumb_encode_method", "temp_path")
    28  
    29  	if f.lastRawExts != ffmpegOpts["thumb_ffmpeg_exts"] {
    30  		f.exts = strings.Split(ffmpegOpts["thumb_ffmpeg_exts"], ",")
    31  	}
    32  
    33  	if !util.IsInExtensionList(f.exts, name) {
    34  		return nil, fmt.Errorf("unsupported video format: %w", ErrPassThrough)
    35  	}
    36  
    37  	tempOutputPath := filepath.Join(
    38  		util.RelativePath(ffmpegOpts["temp_path"]),
    39  		"thumb",
    40  		fmt.Sprintf("thumb_%s.%s", uuid.Must(uuid.NewV4()).String(), ffmpegOpts["thumb_encode_method"]),
    41  	)
    42  
    43  	tempInputPath := src
    44  	if tempInputPath == "" {
    45  		// If not local policy files, download to temp folder
    46  		tempInputPath = filepath.Join(
    47  			util.RelativePath(ffmpegOpts["temp_path"]),
    48  			"thumb",
    49  			fmt.Sprintf("ffmpeg_%s%s", uuid.Must(uuid.NewV4()).String(), filepath.Ext(name)),
    50  		)
    51  
    52  		// Due to limitations of ffmpeg, we need to write the input file to disk first
    53  		tempInputFile, err := util.CreatNestedFile(tempInputPath)
    54  		if err != nil {
    55  			return nil, fmt.Errorf("failed to create temp file: %w", err)
    56  		}
    57  
    58  		defer os.Remove(tempInputPath)
    59  		defer tempInputFile.Close()
    60  
    61  		if _, err = io.Copy(tempInputFile, file); err != nil {
    62  			return nil, fmt.Errorf("failed to write input file: %w", err)
    63  		}
    64  
    65  		tempInputFile.Close()
    66  	}
    67  
    68  	// Invoke ffmpeg
    69  	scaleOpt := fmt.Sprintf("scale=%s:%s:force_original_aspect_ratio=decrease", options["thumb_width"], options["thumb_height"])
    70  	cmd := exec.CommandContext(ctx,
    71  		ffmpegOpts["thumb_ffmpeg_path"], "-ss", ffmpegOpts["thumb_ffmpeg_seek"], "-i", tempInputPath,
    72  		"-vf", scaleOpt, "-vframes", "1", tempOutputPath)
    73  
    74  	// Redirect IO
    75  	var stdErr bytes.Buffer
    76  	cmd.Stdin = file
    77  	cmd.Stderr = &stdErr
    78  
    79  	if err := cmd.Run(); err != nil {
    80  		util.Log().Warning("Failed to invoke ffmpeg: %s", stdErr.String())
    81  		return nil, fmt.Errorf("failed to invoke ffmpeg: %w", err)
    82  	}
    83  
    84  	return &Result{Path: tempOutputPath}, nil
    85  }
    86  
    87  func (f *FfmpegGenerator) Priority() int {
    88  	return 200
    89  }
    90  
    91  func (f *FfmpegGenerator) EnableFlag() string {
    92  	return "thumb_ffmpeg_enabled"
    93  }