From e1a05a057a2c956f665ae980986c879ea20c0bba Mon Sep 17 00:00:00 2001
From: zijiren233 <pyh1670605849@gmail.com>
Date: Thu, 24 Oct 2024 01:45:43 +0800
Subject: [PATCH] fix: pre get m3u8 file in token

---
 server/handlers/movie.go      |  2 +-
 server/handlers/proxy/m3u8.go | 29 +++++++++++++++++++++--------
 utils/m3u8/m3u8.go            | 16 +++++++++-------
 3 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/server/handlers/movie.go b/server/handlers/movie.go
index 1d1e8d3a..c8aa7be1 100644
--- a/server/handlers/movie.go
+++ b/server/handlers/movie.go
@@ -654,7 +654,7 @@ func ServeM3u8(ctx *gin.Context) {
 		ctx.AbortWithStatusJSON(http.StatusBadRequest, model.NewApiErrorStringResp("invalid token"))
 		return
 	}
-	err = proxy.ProxyM3u8(ctx, claims.TargetUrl, m.Movie.MovieBase.Headers, utils.IsM3u8Url(claims.TargetUrl), ctx.GetString("token"), room.ID, m.ID)
+	err = proxy.ProxyM3u8(ctx, claims.TargetUrl, m.Movie.MovieBase.Headers, claims.IsM3u8File, ctx.GetString("token"), room.ID, m.ID)
 	if err != nil {
 		log.Errorf("proxy m3u8 error: %v", err)
 	}
diff --git a/server/handlers/proxy/m3u8.go b/server/handlers/proxy/m3u8.go
index 19ea503d..fdde9e88 100644
--- a/server/handlers/proxy/m3u8.go
+++ b/server/handlers/proxy/m3u8.go
@@ -18,9 +18,10 @@ import (
 )
 
 type m3u8TargetClaims struct {
-	RoomId    string `json:"r"`
-	MovieId   string `json:"m"`
-	TargetUrl string `json:"t"`
+	RoomId     string `json:"r"`
+	MovieId    string `json:"m"`
+	TargetUrl  string `json:"t"`
+	IsM3u8File bool   `json:"f"`
 	jwt.RegisteredClaims
 }
 
@@ -38,11 +39,12 @@ func GetM3u8Target(token string) (*m3u8TargetClaims, error) {
 	return claims, nil
 }
 
-func NewM3u8TargetToken(targetUrl, roomId, movieId string) (string, error) {
+func NewM3u8TargetToken(targetUrl, roomId, movieId string, isM3u8File bool) (string, error) {
 	claims := &m3u8TargetClaims{
-		RoomId:    roomId,
-		MovieId:   movieId,
-		TargetUrl: targetUrl,
+		RoomId:     roomId,
+		MovieId:    movieId,
+		TargetUrl:  targetUrl,
+		IsM3u8File: isM3u8File,
 		RegisteredClaims: jwt.RegisteredClaims{
 			NotBefore: jwt.NewNumericDate(time.Now()),
 		},
@@ -74,8 +76,19 @@ func ProxyM3u8(ctx *gin.Context, u string, headers map[string]string, isM3u8File
 	if err != nil {
 		return fmt.Errorf("read response body error: %w", err)
 	}
+	hasM3u8File := false
+	err = m3u8.RangeM3u8SegmentsWithBaseUrl(stream.BytesToString(b), u, func(segmentUrl string) (bool, error) {
+		if utils.IsM3u8Url(segmentUrl) {
+			hasM3u8File = true
+			return false, nil
+		}
+		return true, nil
+	})
+	if err != nil {
+		return fmt.Errorf("range m3u8 segments with base url error: %w", err)
+	}
 	m3u8Str, err := m3u8.ReplaceM3u8SegmentsWithBaseUrl(stream.BytesToString(b), u, func(segmentUrl string) (string, error) {
-		targetToken, err := NewM3u8TargetToken(segmentUrl, roomId, movieId)
+		targetToken, err := NewM3u8TargetToken(segmentUrl, roomId, movieId, hasM3u8File)
 		if err != nil {
 			return "", err
 		}
diff --git a/utils/m3u8/m3u8.go b/utils/m3u8/m3u8.go
index 30d65e8f..7ed91cb4 100644
--- a/utils/m3u8/m3u8.go
+++ b/utils/m3u8/m3u8.go
@@ -9,9 +9,9 @@ import (
 
 func GetM3u8AllSegments(m3u8Str string, baseUrl string) ([]string, error) {
 	var segments []string
-	err := RangeM3u8SegmentsWithBaseUrl(m3u8Str, baseUrl, func(segmentUrl string) error {
+	err := RangeM3u8SegmentsWithBaseUrl(m3u8Str, baseUrl, func(segmentUrl string) (bool, error) {
 		segments = append(segments, segmentUrl)
-		return nil
+		return true, nil
 	})
 	if err != nil {
 		return nil, err
@@ -19,13 +19,15 @@ func GetM3u8AllSegments(m3u8Str string, baseUrl string) ([]string, error) {
 	return segments, nil
 }
 
-func RangeM3u8Segments(m3u8Str string, callback func(segmentUrl string) error) error {
+func RangeM3u8Segments(m3u8Str string, callback func(segmentUrl string) (bool, error)) error {
 	scanner := bufio.NewScanner(strings.NewReader(m3u8Str))
 	for scanner.Scan() {
 		line := strings.TrimSpace(scanner.Text())
 		if line != "" && !strings.HasPrefix(line, "#") {
-			if err := callback(line); err != nil {
+			if ok, err := callback(line); err != nil {
 				return err
+			} else if !ok {
+				break
 			}
 		}
 	}
@@ -35,16 +37,16 @@ func RangeM3u8Segments(m3u8Str string, callback func(segmentUrl string) error) e
 	return nil
 }
 
-func RangeM3u8SegmentsWithBaseUrl(m3u8Str string, baseUrl string, callback func(segmentUrl string) error) error {
+func RangeM3u8SegmentsWithBaseUrl(m3u8Str string, baseUrl string, callback func(segmentUrl string) (bool, error)) error {
 	baseUrlParsed, err := url.Parse(baseUrl)
 	if err != nil {
 		return fmt.Errorf("parse base url error: %w", err)
 	}
-	return RangeM3u8Segments(m3u8Str, func(segmentUrl string) error {
+	return RangeM3u8Segments(m3u8Str, func(segmentUrl string) (bool, error) {
 		if !strings.HasPrefix(segmentUrl, "http://") && !strings.HasPrefix(segmentUrl, "https://") {
 			segmentUrlParsed, err := url.Parse(segmentUrl)
 			if err != nil {
-				return fmt.Errorf("parse segment url error: %w", err)
+				return false, fmt.Errorf("parse segment url error: %w", err)
 			}
 			segmentUrl = baseUrlParsed.ResolveReference(segmentUrlParsed).String()
 		}