Feat: ts file and ts file type

zijiren233 · Nov 3, 2023 · 6b5ecf1 · 6b5ecf1
1 parent c509743
commit 6b5ecf1
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 19 deletions.
diff --git a/cmd/server.go b/cmd/server.go
@@ -1,8 +1,13 @@
 package cmd
 
 import (
+	"bytes"
 	"fmt"
+	"image"
+	"image/color"
+	"image/png"
 	"log"
+	"math/rand"
 	"net"
 	"net/http"
 	"path"
@@ -50,45 +55,64 @@ func Server(cmd *cobra.Command, args []string) {
 	}
 	e := gin.Default()
 	utils.Cors(e)
-	e.GET("/:app/*channel", func(c *gin.Context) {
-		appName := c.Param("app")
-		channelStr := strings.Trim(c.Param("channel"), "/")
+	e.GET("/:app/*channel", func(ctx *gin.Context) {
+		appName := ctx.Param("app")
+		channelStr := strings.Trim(ctx.Param("channel"), "/")
 		channelSplitd := strings.Split(channelStr, "/")
 		fileName := channelSplitd[0]
 		fileExt := path.Ext(channelStr)
 		channelName := strings.TrimSuffix(fileName, fileExt)
 
 		channel, ok := channels.Load(appName)
 		if !ok {
-			c.AbortWithStatusJSON(http.StatusNotFound, gin.H{
+			ctx.AbortWithStatusJSON(http.StatusNotFound, gin.H{
 				"error": "app not found",
 			})
 			return
 		}
 		switch fileExt {
 		case ".flv":
-			w := httpflv.NewHttpFLVWriter(c.Writer)
+			w := httpflv.NewHttpFLVWriter(ctx.Writer)
 			defer w.Close()
 			channel.AddPlayer(w)
 			w.SendPacket()
 		case ".m3u8":
-			b, err := channel.GenM3U8PlayList(fmt.Sprintf("/%s/%s", appName, channelName))
+			b, err := channel.GenM3U8File(func(tsName string) (tsPath string) {
+				return fmt.Sprintf("/%s/%s/%s.%s", appName, channelName, tsName, ctx.DefaultQuery("t", "ts"))
+			})
 			if err != nil {
-				c.AbortWithStatusJSON(http.StatusNotFound, gin.H{
+				ctx.AbortWithStatusJSON(http.StatusNotFound, gin.H{
 					"error": err.Error(),
 				})
 				return
 			}
-			c.Data(http.StatusOK, hls.M3U8ContentType, b.Bytes())
+			ctx.Data(http.StatusOK, hls.M3U8ContentType, b)
 		case ".ts":
 			b, err := channel.GetTsFile(channelSplitd[1])
 			if err != nil {
-				c.AbortWithStatusJSON(http.StatusNotFound, gin.H{
+				ctx.AbortWithStatusJSON(http.StatusNotFound, gin.H{
+					"error": err.Error(),
+				})
+				return
+			}
+			ctx.Data(http.StatusOK, hls.TSContentType, b)
+		case ".png":
+			b, err := channel.GetTsFile(strings.TrimRight(channelSplitd[1], ".png"))
+			if err != nil {
+				ctx.AbortWithStatusJSON(http.StatusNotFound, gin.H{
 					"error": err.Error(),
 				})
 				return
 			}
-			c.Data(http.StatusOK, hls.TSContentType, b)
+			img := image.NewGray(image.Rect(0, 0, 1, 1))
+			img.Set(1, 1, color.Gray{uint8(rand.Intn(255))})
+			f := bytes.NewBuffer(nil)
+			err = png.Encode(f, img)
+			if err != nil {
+				panic(err)
+			}
+			ctx.Data(http.StatusOK, "image/png", f.Bytes())
+			ctx.Writer.Write(b)
 		}
 	})
 	go http.Serve(httpl, e.Handler())

diff --git a/protocol/hls/cache.go b/protocol/hls/cache.go
@@ -4,7 +4,7 @@ import (
 	"bytes"
 	"fmt"
 	"io/fs"
-	"path"
+	"path/filepath"
 	"strings"
 	"sync"
 
@@ -38,7 +38,7 @@ func (tc *TSCache) all() []*TSItem {
 	return items
 }
 
-func (tc *TSCache) GenM3U8PlayList(tsBasePath string) *bytes.Buffer {
+func (tc *TSCache) GenM3U8File(tsPath func(tsName string) (tsPath string)) []byte {
 	var seq int64
 	var maxDuration int64
 	m3u8body := bytes.NewBuffer(nil)
@@ -53,14 +53,14 @@ func (tc *TSCache) GenM3U8PlayList(tsBasePath string) *bytes.Buffer {
 		if seq == 0 {
 			seq = item.SeqNum
 		}
-		fmt.Fprintf(m3u8body, "#EXTINF:%.3f,\n%s.ts\n#EXT-X-BYTERANGE:%d\n", float64(item.Duration)/float64(1000), path.Join(tsBasePath, item.TsName), len(item.Data))
+		fmt.Fprintf(m3u8body, "#EXTINF:%.3f,\n%s\n#EXT-X-BYTERANGE:%d\n", float64(item.Duration)/float64(1000), tsPath(item.TsName), len(item.Data))
 	}
 	w := bytes.NewBuffer(make([]byte, 0, m3u8body.Len()+256))
 	fmt.Fprintf(w,
 		"#EXTM3U\n#EXT-X-VERSION:3\n#EXT-X-ALLOW-CACHE:NO\n#EXT-X-TARGETDURATION:%d\n#EXT-X-MEDIA-SEQUENCE:%d\n\n",
 		maxDuration/1000+1, seq)
 	m3u8body.WriteTo(w)
-	return w
+	return w.Bytes()
 }
 
 func (tc *TSCache) PushItem(item *TSItem) {
@@ -75,7 +75,7 @@ func (tc *TSCache) PushItem(item *TSItem) {
 }
 
 func (tc *TSCache) GetItem(tsName string) (*TSItem, error) {
-	tsName = strings.TrimSuffix(tsName, ".ts")
+	tsName = strings.TrimSuffix(tsName, filepath.Ext(tsName))
 	tc.lock.RLock()
 	defer tc.lock.RUnlock()
 	for e := tc.l.Front(); e != nil; e = e.Next() {

diff --git a/protocol/hls/source.go b/protocol/hls/source.go
@@ -133,7 +133,7 @@ func (source *Source) cut() {
 		source.flushAudio()
 
 		source.seq++
-		source.tsCache.PushItem(NewTSItem(fmt.Sprint(time.Now().UnixMilli()), source.stat.durationMs(), source.seq, source.btswriter.Bytes()))
+		source.tsCache.PushItem(NewTSItem(fmt.Sprint(time.Now().UnixMicro()), source.stat.durationMs(), source.seq, source.btswriter.Bytes()))
 
 		source.btswriter.Reset()
 		source.stat.resetAndNew()

diff --git a/server/channel.go b/server/channel.go
@@ -1,7 +1,6 @@
 package server
 
 import (
-	"bytes"
 	"errors"
 	"sync"
 	"sync/atomic"
@@ -218,11 +217,11 @@ func (c *Channel) InitdHlsPlayer() bool {
 
 var ErrHlsPlayerNotInit = errors.New("hls player not init")
 
-func (c *Channel) GenM3U8PlayList(tsBashPath string) (*bytes.Buffer, error) {
+func (c *Channel) GenM3U8File(tsPath func(tsName string) (tsPath string)) ([]byte, error) {
 	if !c.InitdHlsPlayer() {
 		return nil, ErrHlsPlayerNotInit
 	}
-	return c.HlsPlayer().GetCacheInc().GenM3U8PlayList(tsBashPath), nil
+	return c.HlsPlayer().GetCacheInc().GenM3U8File(tsPath), nil
 }
 
 func (c *Channel) GetTsFile(tsName string) ([]byte, error) {