tcjlj/MeBox
1
0
Fork 0
mirror of https://github.com/truewhile/MeBox.git synced 2026-10-01 03:56:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: enforce transcoding resource controls

Browse Source
This commit is contained in:
ShukeBta
2026-05-29 13:54:19 +08:00
parent 27df93fa3d
commit a8395c9de7
16 changed files with 363 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/config/config.go
+19 -7
View File
@@ -47,13 +47,19 @@ type ApiConfigConfig struct {
// TranscoderConfig 控制 HLS / ffmpeg 后端。
type TranscoderConfig struct {
Encoder string `mapstructure:"encoder"` // "" / nvenc / qsv / vaapi
Preset string `mapstructure:"preset"`
VideoBitrate string `mapstructure:"video_bitrate"`
MaxRate string `mapstructure:"max_rate"`
BufSize string `mapstructure:"buf_size"`
MaxHeight int `mapstructure:"max_height"`
SegmentSeconds int `mapstructure:"segment_seconds"`
Encoder string `mapstructure:"encoder"` // "" / nvenc / qsv / vaapi
Enabled bool `mapstructure:"enabled"`
HardwareAccel bool `mapstructure:"hardware_accel"`
Preset string `mapstructure:"preset"`
VideoBitrate string `mapstructure:"video_bitrate"`
MaxRate string `mapstructure:"max_rate"`
BufSize string `mapstructure:"buf_size"`
MaxHeight int `mapstructure:"max_height"`
SegmentSeconds int `mapstructure:"segment_seconds"`
Realtime bool `mapstructure:"realtime"`
Threads int `mapstructure:"threads"`
MaxConcurrent int `mapstructure:"max_concurrent"`
IdleTimeoutSeconds int `mapstructure:"idle_timeout_seconds"`
}
// AppConfig 保存运行时应用参数。
@@ -259,12 +265,18 @@ func setDefaults(v *viper.Viper) {
v.SetDefault("organizer.categories.uncategorized_tv", "未分类")
v.SetDefault("transcoder.encoder", "")
v.SetDefault("transcoder.enabled", true)
v.SetDefault("transcoder.hardware_accel", false)
v.SetDefault("transcoder.preset", "veryfast")
v.SetDefault("transcoder.video_bitrate", "1500k")
v.SetDefault("transcoder.max_rate", "1800k")
v.SetDefault("transcoder.buf_size", "3000k")
v.SetDefault("transcoder.max_height", 720)
v.SetDefault("transcoder.segment_seconds", 4)
v.SetDefault("transcoder.realtime", true)
v.SetDefault("transcoder.threads", 2)
v.SetDefault("transcoder.max_concurrent", 1)
v.SetDefault("transcoder.idle_timeout_seconds", 120)
// API Config 默认设置
v.SetDefault("api_config.auto_encrypt", true)
internal/handler/admin.go
+6
View File
@@ -187,6 +187,12 @@ func updateSettingHandler(svc *service.Container) gin.HandlerFunc {
return
}
service.ApplyRuntimeSetting(svc.Cfg, req.Key, req.Value)
if req.Key == "transcode.enabled" && !svc.Cfg.Transcoder.Enabled {
svc.Transcoder.StopAll()
}
if req.Key == "transcode.hw_enabled" || req.Key == "transcode.hw_accel" || req.Key == "transcoder.hardware_accel" || req.Key == "transcoder.encoder" {
svc.Transcoder.StopAll()
}
c.Status(http.StatusNoContent)
}
}
internal/handler/streaming.go
+8
View File
@@ -18,6 +18,14 @@ func hlsPlaylistHandler(svc *service.Container) gin.HandlerFunc {
c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
return
}
if errors.Is(err, service.ErrTranscodeDisabled) {
c.JSON(http.StatusConflict, gin.H{"error": "transcode disabled"})
return
}
if errors.Is(err, service.ErrTranscodeBusy) {
c.JSON(http.StatusTooManyRequests, gin.H{"error": "transcode busy"})
return
}
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
internal/handler/system_extra.go
+5 -1
View File
@@ -64,8 +64,12 @@ func schemaHandler(_ *service.Container) gin.HandlerFunc {
"items": []gin.H{
{"key": "tmdb.language", "type": "select", "label": "TMDb 元数据语言"},
{"key": "transcode.enabled", "type": "toggle", "label": "启用转码"},
{"key": "transcode.hw_accel", "type": "select", "label": "硬件加速"},
{"key": "transcode.hw_accel", "type": "select", "label": "硬件编码器"},
{"key": "transcode.hw_enabled", "type": "toggle", "label": "启用硬件加速"},
{"key": "transcode.max_jobs", "type": "number", "label": "最大并发"},
{"key": "transcode.realtime", "type": "toggle", "label": "按播放速度转码"},
{"key": "transcode.threads", "type": "number", "label": "软件转码线程数"},
{"key": "transcode.idle_timeout_seconds", "type": "number", "label": "转码空闲停止秒数"},
{"key": "ffmpeg.path", "type": "text", "label": "FFmpeg 路径"},
{"key": "ffprobe.path", "type": "text", "label": "FFprobe 路径"},
},
internal/service/runtime_settings.go
+29
View File
@@ -37,6 +37,10 @@ func ApplyRuntimeSetting(cfg *config.Config, key, value string) {
cfg.App.FFmpegPath = value
case "ffprobe.path", "app.ffprobe_path":
cfg.App.FFprobePath = value
case "transcode.enabled", "transcoder.enabled":
cfg.Transcoder.Enabled = parseBoolSetting(value, true)
case "transcode.hw_enabled", "transcoder.hardware_accel":
cfg.Transcoder.HardwareAccel = parseBoolSetting(value, false)
case "transcode.hw_accel", "transcoder.encoder":
switch value {
case "", "auto", "none", "software":
@@ -48,6 +52,20 @@ func ApplyRuntimeSetting(cfg *config.Config, key, value string) {
if n, err := strconv.Atoi(value); err == nil {
cfg.Transcoder.MaxHeight = n
}
case "transcode.max_jobs", "transcoder.max_concurrent":
if n, err := strconv.Atoi(value); err == nil {
cfg.Transcoder.MaxConcurrent = n
}
case "transcode.realtime", "transcoder.realtime":
cfg.Transcoder.Realtime = parseBoolSetting(value, true)
case "transcode.threads", "transcoder.threads":
if n, err := strconv.Atoi(value); err == nil {
cfg.Transcoder.Threads = n
}
case "transcode.idle_timeout_seconds", "transcoder.idle_timeout_seconds":
if n, err := strconv.Atoi(value); err == nil {
cfg.Transcoder.IdleTimeoutSeconds = n
}
case "transcode.video_bitrate", "transcoder.video_bitrate":
cfg.Transcoder.VideoBitrate = value
case "license.server_url":
@@ -56,3 +74,14 @@ func ApplyRuntimeSetting(cfg *config.Config, key, value string) {
cfg.License.HMACSecret = value
}
}
func parseBoolSetting(value string, fallback bool) bool {
switch strings.ToLower(strings.TrimSpace(value)) {
case "1", "true", "yes", "on", "enabled", "启用", "开启":
return true
case "0", "false", "no", "off", "disabled", "禁用", "关闭":
return false
default:
return fallback
}
}
internal/service/runtime_settings_test.go
+33
View File
@@ -0,0 +1,33 @@
package service
import (
"testing"
"github.com/ShukeBta/MediaStationGo/internal/config"
)
func TestApplyRuntimeSettingTranscodeSwitches(t *testing.T) {
cfg := &config.Config{}
cfg.Transcoder.Enabled = true
cfg.Transcoder.HardwareAccel = false
ApplyRuntimeSetting(cfg, "transcode.enabled", "false")
if cfg.Transcoder.Enabled {
t.Fatal("transcode.enabled=false should disable transcoding")
}
ApplyRuntimeSetting(cfg, "transcode.hw_enabled", "true")
if !cfg.Transcoder.HardwareAccel {
t.Fatal("transcode.hw_enabled=true should enable hardware accel")
}
ApplyRuntimeSetting(cfg, "transcode.hw_accel", "nvenc")
if cfg.Transcoder.Encoder != "nvenc" {
t.Fatalf("encoder = %q, want nvenc", cfg.Transcoder.Encoder)
}
ApplyRuntimeSetting(cfg, "transcode.max_jobs", "1")
if cfg.Transcoder.MaxConcurrent != 1 {
t.Fatalf("max concurrent = %d, want 1", cfg.Transcoder.MaxConcurrent)
}
}
internal/service/stream.go
+2
View File
@@ -94,6 +94,7 @@ func (s *StreamService) ServeHLSPlaylist(w http.ResponseWriter, r *http.Request,
if _, err := s.transcoder.EnsureJob(r.Context(), mediaID); err != nil {
return err
}
s.transcoder.TouchJob(mediaID)
if !s.transcoder.WaitReady(r.Context(), mediaID, 30*time.Second) {
return errors.New("hls playlist not ready")
}
@@ -145,6 +146,7 @@ func appendQueryToHLSSegments(playlist, rawQuery string) string {
// ServeHLSSegment writes a single .ts segment from the on-disk cache.
func (s *StreamService) ServeHLSSegment(w http.ResponseWriter, r *http.Request, mediaID, segment string) error {
s.transcoder.TouchJob(mediaID)
// Only allow segments that look like seg_NNNNN.ts so we cannot be tricked
// into reading arbitrary files via path traversal.
if !strings.HasPrefix(segment, "seg_") || !strings.HasSuffix(segment, ".ts") {
internal/service/transcoder.go
+113 -8
View File
@@ -29,6 +29,7 @@ import (
"os"
"os/exec"
"path/filepath"
"strconv"
"strings"
"sync"
"time"
@@ -56,10 +57,19 @@ type hlsJob struct {
outputDir string
cancel context.CancelFunc
startedAt time.Time
lastAccess time.Time
playlistOK bool
encoder string
}
var (
// ErrTranscodeDisabled is returned when HLS transcoding is globally disabled.
ErrTranscodeDisabled = errors.New("transcode disabled")
// ErrTranscodeBusy is returned when the server has reached its configured
// ffmpeg concurrency limit.
ErrTranscodeBusy = errors.New("transcode concurrency limit reached")
)
// NewTranscoderService is the constructor.
func NewTranscoderService(cfg *config.Config, log *zap.Logger, repo *repository.Container, hub *Hub) *TranscoderService {
return &TranscoderService{
@@ -85,6 +95,9 @@ func (t *TranscoderService) PlaylistPath(mediaID string) string {
// non-blocking: it returns the playlist path immediately. The caller is
// expected to poll until WaitReady reports true.
func (t *TranscoderService) EnsureJob(ctx context.Context, mediaID string) (string, error) {
if !t.cfg.Transcoder.Enabled {
return "", ErrTranscodeDisabled
}
m, err := t.repo.Media.FindByID(ctx, mediaID)
if err != nil {
return "", err
@@ -101,9 +114,14 @@ func (t *TranscoderService) EnsureJob(ctx context.Context, mediaID string) (stri
t.mu.Lock()
if _, ok := t.jobs[mediaID]; ok {
t.touchJobLocked(mediaID)
t.mu.Unlock()
return t.PlaylistPath(mediaID), nil
}
if max := t.maxConcurrent(); max > 0 && len(t.jobs) >= max {
t.mu.Unlock()
return "", ErrTranscodeBusy
}
outDir := t.HLSDir(mediaID)
if err := os.MkdirAll(outDir, 0o755); err != nil {
@@ -113,15 +131,17 @@ func (t *TranscoderService) EnsureJob(ctx context.Context, mediaID string) (stri
jobCtx, cancel := context.WithCancel(context.Background())
job := &hlsJob{
mediaID: mediaID,
outputDir: outDir,
cancel: cancel,
startedAt: time.Now(),
encoder: t.cfg.Transcoder.Encoder,
mediaID: mediaID,
outputDir: outDir,
cancel: cancel,
startedAt: time.Now(),
lastAccess: time.Now(),
encoder: t.effectiveEncoder(),
}
t.jobs[mediaID] = job
t.mu.Unlock()
go t.monitorIdle(jobCtx, job)
go t.runFFmpeg(jobCtx, job, m.Path)
return t.PlaylistPath(mediaID), nil
}
@@ -160,6 +180,21 @@ func (t *TranscoderService) StopJob(mediaID string) {
}
}
// TouchJob records client activity for the HLS playlist or segment. The idle
// watchdog uses it to stop ffmpeg soon after the player is closed or switches
// back to direct play.
func (t *TranscoderService) TouchJob(mediaID string) {
t.mu.Lock()
defer t.mu.Unlock()
t.touchJobLocked(mediaID)
}
func (t *TranscoderService) touchJobLocked(mediaID string) {
if j, ok := t.jobs[mediaID]; ok {
j.lastAccess = time.Now()
}
}
// StopAll terminates every running transcode (called on graceful shutdown).
func (t *TranscoderService) StopAll() {
t.mu.Lock()
@@ -194,6 +229,51 @@ func (t *TranscoderService) Active() []ActiveJob {
return out
}
func (t *TranscoderService) maxConcurrent() int {
if t.cfg.Transcoder.MaxConcurrent <= 0 {
return 1
}
return t.cfg.Transcoder.MaxConcurrent
}
func (t *TranscoderService) idleTimeout() time.Duration {
if t.cfg.Transcoder.IdleTimeoutSeconds <= 0 {
return 120 * time.Second
}
return time.Duration(t.cfg.Transcoder.IdleTimeoutSeconds) * time.Second
}
func (t *TranscoderService) monitorIdle(ctx context.Context, job *hlsJob) {
timeout := t.idleTimeout()
ticker := time.NewTicker(15 * time.Second)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
t.mu.Lock()
current, ok := t.jobs[job.mediaID]
if !ok {
t.mu.Unlock()
return
}
idleFor := time.Since(current.lastAccess)
t.mu.Unlock()
if idleFor >= timeout {
t.log.Info("transcode idle timeout",
zap.String("media_id", job.mediaID),
zap.Duration("idle_for", idleFor),
zap.Duration("timeout", timeout),
)
t.StopJob(job.mediaID)
return
}
}
}
}
func (t *TranscoderService) runFFmpeg(ctx context.Context, job *hlsJob, source string) {
bin, err := t.resolveFFmpegPath()
if err != nil {
@@ -249,7 +329,7 @@ func (t *TranscoderService) runFFmpeg(ctx context.Context, job *hlsJob, source s
func (t *TranscoderService) resolveFFmpegPath() (string, error) {
var lastErr error
for _, bin := range executableCandidates(strings.TrimSpace(t.cfg.App.FFmpegPath), "ffmpeg") {
if err := validateFFmpegForTranscode(context.Background(), bin, t.cfg.Transcoder.Encoder); err != nil {
if err := validateFFmpegForTranscode(context.Background(), bin, t.effectiveEncoder()); err != nil {
lastErr = err
continue
}
@@ -262,6 +342,22 @@ func (t *TranscoderService) resolveFFmpegPath() (string, error) {
return "", fmt.Errorf("ffmpeg not found in PATH or common local app directories; configure app.ffmpeg_path to an existing local ffmpeg")
}
func (t *TranscoderService) effectiveEncoder() string {
if !t.cfg.Transcoder.HardwareAccel {
return ""
}
return normalizedHardwareEncoder(t.cfg.Transcoder.Encoder)
}
func normalizedHardwareEncoder(encoder string) string {
switch strings.ToLower(strings.TrimSpace(encoder)) {
case "nvenc", "qsv", "vaapi":
return strings.ToLower(strings.TrimSpace(encoder))
default:
return ""
}
}
func validateFFmpegForTranscode(ctx context.Context, bin, encoder string) error {
required := requiredVideoEncoder(encoder)
out, err := commandOutput(ctx, 8*time.Second, bin, "-hide_banner", "-encoders")
@@ -314,7 +410,10 @@ func hasFFmpegListEntry(output, name string) bool {
// encoder. The function is package-level so the unit test can pin its
// behaviour without spawning a real ffmpeg process.
func buildFFmpegArgs(cfg *config.Config, source, playlist, segments string) []string {
enc := cfg.Transcoder.Encoder
enc := ""
if cfg.Transcoder.HardwareAccel {
enc = normalizedHardwareEncoder(cfg.Transcoder.Encoder)
}
bitrate := cfg.Transcoder.VideoBitrate
if bitrate == "" {
bitrate = "1500k"
@@ -373,11 +472,17 @@ func buildFFmpegArgs(cfg *config.Config, source, playlist, segments string) []st
vpreset = preset
}
args := []string{"-y", "-fflags", "+genpts"}
args := []string{"-y", "-hide_banner", "-nostdin", "-fflags", "+genpts"}
for _, p := range splitNonEmptyArgs(pre) {
args = append(args, p)
}
if cfg.Transcoder.Realtime {
args = append(args, "-re")
}
args = append(args, "-i", source, "-map", "0:v:0?", "-map", "0:a:0?", "-vf", vf, "-c:v", vcodec)
if cfg.Transcoder.Threads > 0 && vcodec == "libx264" {
args = append(args, "-threads", strconv.Itoa(cfg.Transcoder.Threads))
}
if vpreset != "" {
args = append(args, "-preset", vpreset)
}
internal/service/transcoder_test.go
+40 -1
View File
@@ -11,6 +11,8 @@ func TestBuildFFmpegArgs(t *testing.T) {
base := &config.Config{}
base.Transcoder.MaxHeight = 720
base.Transcoder.SegmentSeconds = 4
base.Transcoder.Realtime = true
base.Transcoder.Threads = 2
base.App.VAAPIDevice = "/dev/dri/renderD128"
cases := []struct {
@@ -20,7 +22,7 @@ func TestBuildFFmpegArgs(t *testing.T) {
expectInArgs []string
expectNotPresetIfBlank bool
}{
{"software", "", "libx264", []string{"-preset", "veryfast", "-c:v", "libx264"}, false},
{"software", "", "libx264", []string{"-re", "-preset", "veryfast", "-c:v", "libx264", "-threads", "2"}, false},
{"nvenc", "nvenc", "h264_nvenc", []string{"-hwaccel", "cuda", "-c:v", "h264_nvenc", "-preset", "p4"}, false},
{"qsv", "qsv", "h264_qsv", []string{"-hwaccel", "qsv", "-c:v", "h264_qsv"}, false},
{"vaapi", "vaapi", "h264_vaapi", []string{"-hwaccel", "vaapi", "-vaapi_device", "/dev/dri/renderD128", "-c:v", "h264_vaapi"}, true},
@@ -30,6 +32,7 @@ func TestBuildFFmpegArgs(t *testing.T) {
t.Run(tc.name, func(t *testing.T) {
cfg := *base
cfg.Transcoder.Encoder = tc.encoder
cfg.Transcoder.HardwareAccel = tc.encoder != ""
args := buildFFmpegArgs(&cfg, "/x.mkv", "/o/x.m3u8", "/o/seg_%05d.ts")
joined := strings.Join(args, " ")
for _, frag := range tc.expectInArgs {
@@ -45,6 +48,42 @@ func TestBuildFFmpegArgs(t *testing.T) {
}
}
func TestBuildFFmpegArgsIgnoresEncoderWhenHardwareAccelDisabled(t *testing.T) {
cfg := &config.Config{}
cfg.Transcoder.Encoder = "nvenc"
cfg.Transcoder.HardwareAccel = false
cfg.Transcoder.MaxHeight = 720
cfg.Transcoder.SegmentSeconds = 4
cfg.Transcoder.Realtime = true
cfg.Transcoder.Threads = 2
args := buildFFmpegArgs(cfg, "/x.mkv", "/o/x.m3u8", "/o/seg_%05d.ts")
joined := strings.Join(args, " ")
if strings.Contains(joined, "h264_nvenc") {
t.Fatalf("hardware disabled should not use nvenc, got: %s", joined)
}
if !strings.Contains(joined, "libx264") {
t.Fatalf("hardware disabled should fall back to libx264, got: %s", joined)
}
}
func TestBuildFFmpegArgsCanDisableRealtimeAndThreadCap(t *testing.T) {
cfg := &config.Config{}
cfg.Transcoder.MaxHeight = 720
cfg.Transcoder.SegmentSeconds = 4
cfg.Transcoder.Realtime = false
cfg.Transcoder.Threads = 0
args := buildFFmpegArgs(cfg, "/x.mkv", "/o/x.m3u8", "/o/seg_%05d.ts")
joined := " " + strings.Join(args, " ") + " "
if strings.Contains(joined, " -re ") {
t.Fatalf("realtime=false should not include -re, got: %s", joined)
}
if strings.Contains(joined, " -threads ") {
t.Fatalf("threads=0 should not include -threads, got: %s", joined)
}
}
func TestRequiredVideoEncoder(t *testing.T) {
cases := map[string]string{
"": "libx264",