From 7e05c32d5ba9df39f4fdd737a1e6f668e591dd8c Mon Sep 17 00:00:00 2001
From: zijiren233 <pyh1670605849@gmail.com>
Date: Wed, 23 Oct 2024 23:40:34 +0800
Subject: [PATCH] feat: m3u8 live proxy

---
 internal/op/movie.go     |  3 +++
 server/handlers/movie.go | 34 +++++++++++++++++++++++-----------
 utils/utils.go           |  4 ++++
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/internal/op/movie.go b/internal/op/movie.go
index b68cf11..65e07ef 100644
--- a/internal/op/movie.go
+++ b/internal/op/movie.go
@@ -201,6 +201,9 @@ func (m *Movie) initChannel() (*rtmps.Channel, error) {
 			}()
 			return c, nil
 		case "http", "https":
+			if utils.IsM3u8Url(m.Movie.MovieBase.Url) {
+				return nil, errors.New("m3u8 url not support")
+			}
 			c, init := m.compareAndSwapInitChannel()
 			if !init {
 				return c, nil
diff --git a/server/handlers/movie.go b/server/handlers/movie.go
index 83a4dfe..83660e2 100644
--- a/server/handlers/movie.go
+++ b/server/handlers/movie.go
@@ -72,7 +72,7 @@ func genMovieInfo(
 		if err != nil {
 			return nil, err
 		}
-	} else if movie.MovieBase.RtmpSource || movie.MovieBase.Live && movie.MovieBase.Proxy {
+	} else if movie.MovieBase.RtmpSource {
 		movie.MovieBase.Url = fmt.Sprintf("/api/room/movie/live/hls/list/%s.m3u8?token=%s&roomId=%s", movie.ID, userToken, opMovie.RoomID)
 		movie.MovieBase.Type = "m3u8"
 		movie.MoreSources = append(movie.MoreSources, &dbModel.MoreSource{
@@ -81,6 +81,17 @@ func genMovieInfo(
 			Type: "flv",
 		})
 		movie.MovieBase.Headers = nil
+	} else if movie.MovieBase.Live && movie.MovieBase.Proxy {
+		if !utils.IsM3u8Url(movie.MovieBase.Url) {
+			movie.MoreSources = append(movie.MoreSources, &dbModel.MoreSource{
+				Name: "flv",
+				Url:  fmt.Sprintf("/api/room/movie/live/flv/%s.flv?token=%s&roomId=%s", movie.ID, userToken, opMovie.RoomID),
+				Type: "flv",
+			})
+		}
+		movie.MovieBase.Url = fmt.Sprintf("/api/room/movie/live/hls/list/%s.m3u8?token=%s&roomId=%s", movie.ID, userToken, opMovie.RoomID)
+		movie.MovieBase.Type = "m3u8"
+		movie.MovieBase.Headers = nil
 	} else if movie.MovieBase.Proxy {
 		movie.MovieBase.Url = fmt.Sprintf("/api/room/movie/proxy/%s?token=%s&roomId=%s", movie.ID, userToken, opMovie.RoomID)
 		movie.MovieBase.Headers = nil
@@ -595,7 +606,7 @@ func ProxyMovie(ctx *gin.Context) {
 		// TODO: cache mpd file
 		fallthrough
 	default:
-		if strings.HasPrefix(m.Movie.MovieBase.Type, "m3u") || isM3u8Url(m.Movie.MovieBase.Url) {
+		if strings.HasPrefix(m.Movie.MovieBase.Type, "m3u") || utils.IsM3u8Url(m.Movie.MovieBase.Url) {
 			err = proxyM3u8(ctx, m.Movie.MovieBase.Url, m.Movie.MovieBase.Headers, true, ctx.GetString("token"), room.ID, m.ID)
 			if err != nil {
 				log.Errorf("proxy movie error: %v", err)
@@ -610,10 +621,6 @@ func ProxyMovie(ctx *gin.Context) {
 	}
 }
 
-func isM3u8Url(u string) bool {
-	return strings.HasPrefix(utils.GetUrlExtension(u), "m3u")
-}
-
 type m3u8TargetClaims struct {
 	RoomId    string `json:"r"`
 	MovieId   string `json:"m"`
@@ -710,13 +717,13 @@ func ServeM3u8(ctx *gin.Context) {
 		return
 	}
 
-	if !m.Movie.MovieBase.Proxy {
-		ctx.AbortWithStatusJSON(http.StatusBadRequest, model.NewApiErrorStringResp("proxy is not enabled"))
+	if m.Movie.MovieBase.RtmpSource {
+		ctx.AbortWithStatusJSON(http.StatusBadRequest, model.NewApiErrorStringResp("this movie is rtmp source, not support use this method proxy"))
 		return
 	}
 
-	if m.Movie.MovieBase.Live || m.Movie.MovieBase.RtmpSource {
-		ctx.AbortWithStatusJSON(http.StatusBadRequest, model.NewApiErrorStringResp("this movie is live or rtmp source, not support use this method proxy"))
+	if !m.Movie.MovieBase.Proxy {
+		ctx.AbortWithStatusJSON(http.StatusBadRequest, model.NewApiErrorStringResp("proxy is not enabled"))
 		return
 	}
 
@@ -731,7 +738,7 @@ func ServeM3u8(ctx *gin.Context) {
 		ctx.AbortWithStatusJSON(http.StatusBadRequest, model.NewApiErrorStringResp("invalid token"))
 		return
 	}
-	err = proxyM3u8(ctx, claims.TargetUrl, m.Movie.MovieBase.Headers, isM3u8Url(claims.TargetUrl), ctx.GetString("token"), room.ID, m.ID)
+	err = proxyM3u8(ctx, claims.TargetUrl, m.Movie.MovieBase.Headers, utils.IsM3u8Url(claims.TargetUrl), ctx.GetString("token"), room.ID, m.ID)
 	if err != nil {
 		log.Errorf("proxy m3u8 error: %v", err)
 	}
@@ -860,6 +867,11 @@ func JoinHlsLive(ctx *gin.Context) {
 		ctx.AbortWithStatusJSON(http.StatusBadRequest, model.NewApiErrorStringResp("live proxy is not enabled"))
 		return
 	}
+
+	if utils.IsM3u8Url(m.Movie.MovieBase.Url) {
+		_ = proxyM3u8(ctx, m.Movie.MovieBase.Url, m.Movie.MovieBase.Headers, true, ctx.GetString("token"), room.ID, m.ID)
+		return
+	}
 	channel, err := m.Channel()
 	if err != nil {
 		log.Errorf("join hls live error: %v", err)
diff --git a/utils/utils.go b/utils/utils.go
index 603ce0d..377d85d 100644
--- a/utils/utils.go
+++ b/utils/utils.go
@@ -352,6 +352,10 @@ func GetUrlExtension(u string) string {
 	return GetFileExtension(p.Path)
 }
 
+func IsM3u8Url(u string) bool {
+	return strings.HasPrefix(GetUrlExtension(u), "m3u")
+}
+
 var (
 	needColor     bool
 	needColorOnce sync.Once