From 5ee1de64a3bc3a205855d394701fba90a245869c Mon Sep 17 00:00:00 2001
From: banshan <1310049536@qq.com>
Date: Sun, 27 Jun 2021 14:42:39 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E6=96=87=E4=BB=B6=E6=8A=BD?=
 =?UTF-8?q?=E8=B1=A1=E5=92=8C=E6=96=87=E4=BB=B6=E6=B5=81=E8=BF=87=E6=BB=A4?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 flv.go  | 20 +++++++++++++-------
 main.go | 15 ++++++++++++++-
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/flv.go b/flv.go
index 694418c..1380182 100644
--- a/flv.go
+++ b/flv.go
@@ -10,7 +10,7 @@ import (
 	"github.com/Monibuca/utils/v3/codec"
 )
 
-func getDuration(file *os.File) uint32 {
+func getDuration(file FileWr) uint32 {
 	_, err := file.Seek(-4, io.SeekEnd)
 	if err == nil {
 		var tagSize uint32
@@ -34,12 +34,18 @@ func SaveFlv(streamPath string, append bool) error {
 		flag = flag | os.O_TRUNC | os.O_WRONLY
 	}
 	filePath := filepath.Join(config.Path, streamPath+".flv")
-	if err := os.MkdirAll(filepath.Dir(filePath), 0755); err != nil {
-		return err
-	}
-	file, err := os.OpenFile(filePath, flag, 0755)
-	if err != nil {
-		return err
+	var file FileWr
+	var err error
+	if ExtraConfig.CreateFileFn == nil {
+		if err := os.MkdirAll(filepath.Dir(filePath), 0755); err != nil {
+			return err
+		}
+		file, err = os.OpenFile(filePath, flag, 0755)
+		if err != nil {
+			return err
+		}
+	}else {
+		file,err  =  ExtraConfig.CreateFileFn(filePath)
 	}
 	// return avformat.WriteFLVTag(file, packet)
 	p := Subscriber{
diff --git a/main.go b/main.go
index 05b36ba..6348cf4 100644
--- a/main.go
+++ b/main.go
@@ -2,6 +2,7 @@ package record
 
 import (
 	"encoding/json"
+	"io"
 	"net/http"
 	"os"
 	"path"
@@ -25,6 +26,18 @@ type FlvFileInfo struct {
 	Duration uint32
 }
 
+type FileWr interface {
+	io.Reader
+	io.Writer
+	io.Seeker
+	io.Closer
+}
+
+var ExtraConfig struct {
+	CreateFileFn func(filename string) (FileWr,error)
+	AutoRecordFilter func(stream string) bool
+}
+
 func init() {
 	InstallPlugin(&PluginConfig{
 		Name:   "Record",
@@ -115,7 +128,7 @@ func run() {
 
 func onPublish(v interface{}) {
 	p := v.(*Stream)
-	if config.AutoRecord {
+	if config.AutoRecord || (ExtraConfig.AutoRecordFilter != nil && ExtraConfig.AutoRecordFilter(p.StreamPath)) {
 		SaveFlv(p.StreamPath, false)
 	}
 }