Support m3u8.Master URL and continue download task

hlsdl.go

@@ -1,6 +1,8 @@
 package hlsdl
 
 import (
+	"crypto/md5"
+	"encoding/hex"
 	"errors"
 	"fmt"
 	"io"
@@ -77,15 +79,20 @@ func wait(wg *sync.WaitGroup) chan bool {
 }
 
 func (hlsDl *HlsDl) downloadSegment(segment *Segment) error {
+	if _, err := os.Stat(segment.Path); err == nil {
+		return nil
+	}
+	tmpfile := segment.Path + ".tmp"
 	hlsDl.client.SetRetryCount(5).SetRetryWaitTime(time.Second)
-	resp, err := hlsDl.client.R().SetHeaders(hlsDl.headers).SetOutput(segment.Path).Get(segment.URI)
+	resp, err := hlsDl.client.R().SetHeaders(hlsDl.headers).SetOutput(tmpfile).Get(segment.URI)
 	if err != nil {
 		return err
 	}
 	if resp.StatusCode() != http.StatusOK {
 		return errors.New(resp.Status())
 	}
-	return nil
+	err = os.Rename(tmpfile, segment.Path)
+	return err
 }
 
 func (hlsDl *HlsDl) downloadSegments(segmentsDir string, segments []*Segment) error {
@@ -202,10 +209,13 @@ func (hlsDl *HlsDl) Download() (string, error) {
 	if err != nil {
 		return "", err
 	}
-	segmentsDir := filepath.Join(hlsDl.dir, fmt.Sprintf("%d", hlsDl.startTime))
+
+	tmpdir := md5.Sum([]byte(hlsDl.hlsURL))
+	segmentsDir := filepath.Join(hlsDl.dir, hex.EncodeToString(tmpdir[:]))
 	if err := os.MkdirAll(segmentsDir, os.ModePerm); err != nil {
 		return "", err
 	}
+	log.Println("Tmpdir", segmentsDir)
 	if err := hlsDl.downloadSegments(segmentsDir, segs); err != nil {
 		return "", err
 	}

parser.go

@@ -17,7 +17,7 @@ func parseHlsSegments(hlsURL string, headers map[string]string) ([]*Segment, err
 		return nil, errors.New("Invalid m3u8 url")
 	}
 
-	p, t, err := getM3u8ListType(hlsURL, headers)
+	p, t, err := getM3u8ListType(hlsURL, headers, 0)
 	if err != nil {
 		return nil, err
 	}
@@ -61,10 +61,10 @@ func parseHlsSegments(hlsURL string, headers map[string]string) ([]*Segment, err
 	return segments, nil
 }
 
-func getM3u8ListType(url string, headers map[string]string) (m3u8.Playlist, m3u8.ListType, error) {
+func getM3u8ListType(hlsurl string, headers map[string]string, depth int) (m3u8.Playlist, m3u8.ListType, error) {
 	client := resty.New()
 	client.SetRetryCount(5).SetRetryWaitTime(time.Second)
-	resp, err := client.R().SetHeaders(headers).Get(url)
+	resp, err := client.R().SetHeaders(headers).Get(hlsurl)
 	if err != nil {
 		return nil, 0, err
 	}
@@ -75,5 +75,20 @@ func getM3u8ListType(url string, headers map[string]string) (m3u8.Playlist, m3u8
 	if err != nil {
 		return nil, 0, err
 	}
+	if t == m3u8.MASTER {
+		pt := p.(*m3u8.MasterPlaylist)
+		for _, pv := range pt.Variants {
+			var uo, err = url.Parse(pv.URI)
+			if err != nil {
+				return nil, 0, err
+			}
+			if !uo.IsAbs() {
+				var uo2, _ = url.Parse(hlsurl)
+				uo = uo2.ResolveReference(uo)
+			}
+			var ut = uo.String()
+			return getM3u8ListType(ut, headers, depth+1)
+		}
+	}
 	return p, t, nil
 }

puller.go

@@ -31,7 +31,7 @@ func pullSegment(hlsURL string, quitSignal chan os.Signal) chan *SegmentPuller {
 		pulledSegment := map[uint64]bool{}
 
 		for {
-			p, t, err := getM3u8ListType(hlsURL, nil)
+			p, t, err := getM3u8ListType(hlsURL, nil, 0)
 			if err != nil {
 				c <- &SegmentPuller{Err: err}
 				return