mirror of
https://github.com/langhuihui/monibuca.git
synced 2025-09-27 14:22:08 +08:00
98 lines
2.6 KiB
Go
98 lines
2.6 KiB
Go
package plugin_test
|
|
|
|
import (
|
|
"fmt"
|
|
"os/exec"
|
|
"reflect"
|
|
"strings"
|
|
|
|
task "github.com/langhuihui/gotask"
|
|
"m7s.live/v5"
|
|
"m7s.live/v5/pkg"
|
|
"m7s.live/v5/pkg/format"
|
|
"m7s.live/v5/pkg/util"
|
|
)
|
|
|
|
var ffmpegArgs = []string{
|
|
"-hide_banner",
|
|
"-i", "pipe:0",
|
|
"-vframes", "1",
|
|
"-f", "mjpeg", "pipe:1",
|
|
}
|
|
|
|
type SnapshotTask struct {
|
|
TestBaseTask
|
|
}
|
|
|
|
func (st *SnapshotTask) Run() error {
|
|
var cmd *exec.Cmd
|
|
streamPath := st.StreamPath
|
|
if st.Format == "" {
|
|
// 创建订阅配置
|
|
subConfig := st.testCase.Plugin.GetCommonConf().Subscribe
|
|
subConfig.SubType = m7s.SubscribeTypeTransform
|
|
subConfig.IFrameOnly = true
|
|
|
|
subscriber, err := st.testCase.Plugin.SubscribeWithConfig(st, streamPath, subConfig)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to subscribe to stream: %w", err)
|
|
}
|
|
var annexB *format.AnnexB
|
|
track := subscriber.Publisher.GetVideoTrack(reflect.TypeOf(annexB))
|
|
track.WaitReady()
|
|
reader := pkg.NewAVRingReader(track, "annexb")
|
|
err = reader.ReadFrame(&subConfig)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to read frame: %w", err)
|
|
}
|
|
annexB = reader.Value.Wraps[track.WrapIndex].(*format.AnnexB)
|
|
var mem util.Memory
|
|
mem.CopyFrom(&annexB.Memory)
|
|
reader.StopRead()
|
|
cmd = exec.CommandContext(st, "ffmpeg", ffmpegArgs...)
|
|
r := mem.NewReader()
|
|
cmd.Stdin = &r
|
|
} else {
|
|
cmd = exec.CommandContext(st, "ffmpeg", "-hide_banner", "-i", st.GetInput(streamPath), "-vframes", "1", "-f", "mjpeg", "pipe:1")
|
|
}
|
|
var buf util.Buffer
|
|
cmd.Stderr = st
|
|
cmd.Stdout = &buf
|
|
st.Info("starting ffmpeg", "args", strings.Join(cmd.Args, " "))
|
|
err := cmd.Start()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to start ffmpeg: %w", err)
|
|
}
|
|
st.OnStop(cmd.Process.Kill)
|
|
cmd.Wait()
|
|
if buf.Len() == 0 {
|
|
return fmt.Errorf("snapshot output is empty")
|
|
}
|
|
st.Info("Snapshot completed successfully", "outputSize", buf.Len())
|
|
return task.ErrTaskComplete
|
|
}
|
|
|
|
func (st *SnapshotTask) GetInput(streamPath string) string {
|
|
switch st.Format {
|
|
case "rtmp", "rtsp":
|
|
return fmt.Sprintf("%s://%s/%s", st.Format, st.ServerAddr, streamPath)
|
|
case "srt":
|
|
return fmt.Sprintf("srt://%s?streamid=subscribe:/%s", st.ServerAddr, streamPath)
|
|
case "flv":
|
|
return fmt.Sprintf("http://%s/flv/%s.flv", st.ServerAddr, streamPath)
|
|
case "mp4":
|
|
return fmt.Sprintf("http://%s/mp4/%s.mp4", st.ServerAddr, streamPath)
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func (st *SnapshotTask) Write(buf []byte) (int, error) {
|
|
return st.testCase.Write(append([]byte("[snapshot]"), buf...))
|
|
}
|
|
|
|
func init() {
|
|
testTaskFactory.Register("snapshot", func(s *TestCase, conf TestTaskConfig) task.ITask {
|
|
return &SnapshotTask{TestBaseTask: TestBaseTask{testCase: s, TestTaskConfig: conf}}
|
|
})
|
|
}
|