From fb51dc781d0d9b93299a1b04cf7805ef399d6e6c Mon Sep 17 00:00:00 2001
From: Alexey Khit <alexey.khit@gmail.com>
Date: Sun, 20 Aug 2023 16:35:09 +0300
Subject: [PATCH] Improve HLS reader

---
 pkg/hls/reader.go | 96 ++++++++++++++++++++++++++---------------------
 1 file changed, 53 insertions(+), 43 deletions(-)

diff --git a/pkg/hls/reader.go b/pkg/hls/reader.go
index a0e79b9e..807bd6cf 100644
--- a/pkg/hls/reader.go
+++ b/pkg/hls/reader.go
@@ -2,7 +2,6 @@ package hls
 
 import (
 	"bytes"
-	"errors"
 	"io"
 	"net/http"
 	"net/url"
@@ -16,6 +15,7 @@ type reader struct {
 	client  *http.Client
 	request *http.Request
 
+	playlist    []byte
 	lastSegment []byte
 	lastTime    time.Time
 
@@ -28,18 +28,22 @@ func NewReader(u *url.URL, body io.ReadCloser) (io.Reader, error) {
 		return nil, err
 	}
 
+	var rawURL string
+
 	re := regexp.MustCompile(`#EXT-X-STREAM-INF.+?\n(\S+)`)
 	m := re.FindSubmatch(b)
-	if m == nil {
-		return nil, errors.New("hls: wrong playlist: " + string(b))
+	if m != nil {
+		ref, err := url.Parse(string(m[1]))
+		if err != nil {
+			return nil, err
+		}
+
+		rawURL = u.ResolveReference(ref).String()
+	} else {
+		rawURL = u.String()
 	}
 
-	ref, err := url.Parse(string(m[1]))
-	if err != nil {
-		return nil, err
-	}
-
-	req, err := http.NewRequest("GET", u.ResolveReference(ref).String(), nil)
+	req, err := http.NewRequest("GET", rawURL, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -75,48 +79,52 @@ func (r *reader) Read(dst []byte) (n int, err error) {
 }
 
 func (r *reader) getSegment() ([]byte, error) {
-	for {
-		if wait := time.Second - time.Since(r.lastTime); wait > 0 {
-			time.Sleep(wait)
+	for i := 0; i < 5; i++ {
+		if r.playlist == nil {
+			if wait := time.Second - time.Since(r.lastTime); wait > 0 {
+				time.Sleep(wait)
+			}
+
+			// 1. Load playlist
+			res, err := r.client.Do(r.request)
+			if err != nil {
+				return nil, err
+			}
+
+			r.playlist, err = io.ReadAll(res.Body)
+			if err != nil {
+				return nil, err
+			}
+
+			r.lastTime = time.Now()
+
+			//log.Printf("[hls] load playlist\n%s", r.playlist)
 		}
 
-		// 1. Load playlist
-		res, err := r.client.Do(r.request)
-		if err != nil {
-			return nil, err
-		}
+		for r.playlist != nil {
+			// 2. Remove all previous segments from playlist
+			if i := bytes.Index(r.playlist, r.lastSegment); i > 0 {
+				r.playlist = r.playlist[i:]
+			}
 
-		playlist, err := io.ReadAll(res.Body)
-		if err != nil {
-			return nil, err
-		}
-
-		r.lastTime = time.Now()
-
-		//log.Printf("[hls] load playlist\n%s", playlist)
-
-		// 2. Remove all previous segments from playlist
-		if i := bytes.Index(playlist, r.lastSegment); i > 0 {
-			playlist = playlist[i:]
-		}
-
-		for playlist != nil {
 			// 3. Get link to new segment
-			var segment []byte
-			if segment, playlist = getSegment(playlist); segment == nil {
+			segment := getSegment(r.playlist)
+			if segment == nil {
+				r.playlist = nil
 				break
 			}
 
 			//log.Printf("[hls] load segment: %s", segment)
 
-			ref, err2 := url.Parse(string(segment))
-			if err2 != nil {
-				return nil, err2
+			ref, err := url.Parse(string(segment))
+			if err != nil {
+				return nil, err
 			}
 
 			ref = r.request.URL.ResolveReference(ref)
-			if res, err2 = r.client.Get(ref.String()); err2 != nil {
-				return nil, err2
+			res, err := r.client.Get(ref.String())
+			if err != nil {
+				return nil, err
 			}
 
 			r.lastSegment = segment
@@ -124,23 +132,25 @@ func (r *reader) getSegment() ([]byte, error) {
 			return io.ReadAll(res.Body)
 		}
 	}
+
+	return nil, io.EOF
 }
 
-func getSegment(src []byte) (segment, left []byte) {
+func getSegment(src []byte) []byte {
 	for ok := false; !ok; {
 		ok = bytes.HasPrefix(src, []byte("#EXTINF"))
 
 		i := bytes.IndexByte(src, '\n') + 1
 		if i == 0 {
-			return nil, nil
+			return nil
 		}
 
 		src = src[i:]
 	}
 
 	if i := bytes.IndexByte(src, '\n'); i > 0 {
-		return src[:i], src[i+1:]
+		return src[:i]
 	}
 
-	return src, nil
+	return src
 }