optimize code

downloader/downloader.go

@@ -7,9 +7,11 @@ import (
 	"io"
 	"net/http"
 	"os"
+	"sync"
 	"time"
 
 	"github.com/cheggaaa/pb"
+
 	"github.com/iawia002/annie/config"
 	"github.com/iawia002/annie/request"
 	"github.com/iawia002/annie/utils"
@@ -29,7 +31,7 @@ func Caption(url, refer, fileName, ext string) error {
 		return nil
 	}
 	fmt.Println("\nDownloading captions...")
-	body, err := request.Get(url, refer, nil)
+	body, err := request.GetByte(url, refer, nil)
 	if err != nil {
 		return err
 	}
@@ -42,23 +44,27 @@ func Caption(url, refer, fileName, ext string) error {
 		return fileError
 	}
 	defer file.Close()
-	file.WriteString(body)
+
+	if _, err = file.Write(body); err != nil {
+		return err
+	}
 	return nil
 }
 
 func writeFile(
 	url string, file *os.File, headers map[string]string, bar *pb.ProgressBar,
 ) (int64, error) {
-	res, err := request.Request("GET", url, nil, headers)
+	res, err := request.Request(http.MethodGet, url, nil, headers)
 	if err != nil {
 		return 0, err
 	}
 	defer res.Body.Close()
+
 	writer := io.MultiWriter(file, bar)
 	// Note that io.Copy reads 32kb(maximum) from input and writes them to output, then repeats.
 	// So don't worry about memory.
 	written, copyErr := io.Copy(writer, res.Body)
-	if copyErr != nil {
+	if copyErr != nil && copyErr != io.EOF {
 		return written, fmt.Errorf("file copy error: %s", copyErr)
 	}
 	return written, nil
@@ -110,6 +116,17 @@ func Save(
 	if fileError != nil {
 		return fileError
 	}
+
+	// close and rename temp file at the end of this function
+	defer func() {
+		// must close the file before rename or it will cause
+		// `The process cannot access the file because it is being used by another process.` error.
+		file.Close()
+		if err == nil {
+			os.Rename(tempFilePath, filePath)
+		}
+	}()
+
 	if chunkSizeMB > 0 {
 		var start, end, chunkSize int64
 		chunkSize = int64(chunkSizeMB) * 1024 * 1024
@@ -155,15 +172,6 @@ func Save(
 		}
 	}
 
-	// close and rename temp file at the end of this function
-	defer func() {
-		// must close the file before rename or it will cause
-		// `The process cannot access the file because it is being used by another process.` error.
-		file.Close()
-		if err == nil {
-			os.Rename(tempFilePath, filePath)
-		}
-	}()
 	return nil
 }
 
@@ -220,20 +228,25 @@ func Download(v Data, refer string, chunkSizeMB int) error {
 				return err
 			}
 			reqURL := fmt.Sprintf("%s://%s/jsonrpc", config.Aria2Method, config.Aria2Addr)
-			req, err := http.NewRequest("POST", reqURL, bytes.NewBuffer(jsonData))
+			req, err := http.NewRequest(http.MethodPost, reqURL, bytes.NewBuffer(jsonData))
 			if err != nil {
 				return err
 			}
 			req.Header.Set("Content-Type", "application/json")
-			var client http.Client
-			_, err = client.Do(req)
+
+			var client = http.Client{Timeout: 30 * time.Second}
+			res, err := client.Do(req)
 			if err != nil {
 				return err
 			}
+			// The http Client and Transport guarantee that Body is always
+			// non-nil, even on responses without a body or responses with
+			// a zero-length body.
+			res.Body.Close()
 		}
 		return nil
 	}
-	var err error
+
 	// Skip the complete file that has been merged
 	mergedFilePath, err := utils.FilePath(title, "mp4", false)
 	if err != nil {
@@ -262,8 +275,13 @@ func Download(v Data, refer string, chunkSizeMB int) error {
 	wgp := utils.NewWaitGroupPool(config.ThreadNumber)
 	// multiple fragments
 	errs := make([]error, 0)
+	lock := sync.Mutex{}
 	parts := make([]string, len(data.URLs))
 	for index, url := range data.URLs {
+		if len(errs) > 0 {
+			break
+		}
+
 		partFileName := fmt.Sprintf("%s[%d]", title, index)
 		partFilePath, err := utils.FilePath(partFileName, url.Ext, false)
 		if err != nil {
@@ -276,7 +294,9 @@ func Download(v Data, refer string, chunkSizeMB int) error {
 			defer wgp.Done()
 			err := Save(url, refer, fileName, bar, chunkSizeMB)
 			if err != nil {
+				lock.Lock()
 				errs = append(errs, err)
+				lock.Unlock()
 			}
 		}(url, refer, partFileName, bar)
 	}
@@ -296,8 +316,6 @@ func Download(v Data, refer string, chunkSizeMB int) error {
 	} else {
 		err = utils.MergeToMP4(parts, mergedFilePath, title)
 	}
-	if err != nil {
-		return err
-	}
-	return nil
+
+	return err
 }

downloader/types.go

@@ -47,9 +47,6 @@ type Data struct {
 	URL string `json:"url"`
 }
 
-// EmptyList empty Data list
-var EmptyList = make([]Data, 0)
-
 // EmptyData returns an "empty" Data object with the given URL and error
 func EmptyData(url string, err error) Data {
 	return Data{

extractors/bcy/bcy.go

@@ -2,9 +2,11 @@ package bcy
 
 import (
 	"encoding/json"
+	"fmt"
 	"strings"
 
 	"github.com/iawia002/annie/downloader"
+	"github.com/iawia002/annie/extractors"
 	"github.com/iawia002/annie/parser"
 	"github.com/iawia002/annie/request"
 	"github.com/iawia002/annie/utils"
@@ -22,21 +24,27 @@ type bcyData struct {
 
 // Extract is the main function for extracting data
 func Extract(url string) ([]downloader.Data, error) {
-	var err error
 	html, err := request.Get(url, url, nil)
 	if err != nil {
-		return downloader.EmptyList, err
+		return nil, err
 	}
 
 	// parse json data
 	rep := strings.NewReplacer(`\"`, `"`, `\\`, `\`)
-	jsonString := rep.Replace(utils.MatchOneOf(html, `JSON.parse\("(.+?)"\);`)[1])
+	realURLs := utils.MatchOneOf(html, `JSON.parse\("(.+?)"\);`)
+	if realURLs == nil || len(realURLs) < 2 {
+		return nil, extractors.ErrURLParseFailed
+	}
+	jsonString := rep.Replace(realURLs[1])
+
 	var data bcyData
-	json.Unmarshal([]byte(jsonString), &data)
+	if err = json.Unmarshal([]byte(jsonString), &data); err != nil {
+		return nil, fmt.Errorf("json unmarshal failed, err: %v", err)
+	}
 
 	doc, err := parser.GetDoc(html)
 	if err != nil {
-		return downloader.EmptyList, err
+		return nil, err
 	}
 	title := strings.Replace(parser.Title(doc), " - 半次元 banciyuan - ACG爱好者社区", "", -1)
 
@@ -45,12 +53,12 @@ func Extract(url string) ([]downloader.Data, error) {
 	for _, img := range data.Detail.PostData.Multi {
 		size, err := request.Size(img.OriginalPath, url)
 		if err != nil {
-			return downloader.EmptyList, err
+			return nil, err
 		}
 		totalSize += size
 		_, ext, err := utils.GetNameAndExt(img.OriginalPath)
 		if err != nil {
-			return downloader.EmptyList, err
+			return nil, err
 		}
 		urls = append(urls, downloader.URL{
 			URL:  img.OriginalPath,

extractors/bilibili/bilibili.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/iawia002/annie/config"
 	"github.com/iawia002/annie/downloader"
+	"github.com/iawia002/annie/extractors"
 	"github.com/iawia002/annie/parser"
 	"github.com/iawia002/annie/request"
 	"github.com/iawia002/annie/utils"
@@ -111,7 +112,7 @@ func extractBangumi(url, html string) ([]downloader.Data, error) {
 	var data bangumiData
 	err := json.Unmarshal([]byte(dataString), &data)
 	if err != nil {
-		return downloader.EmptyList, err
+		return nil, err
 	}
 	if !config.Playlist {
 		options := bilibiliOptions{
@@ -173,7 +174,7 @@ func getMultiPageData(html string) (*multiPage, error) {
 func extractNormalVideo(url, html string) ([]downloader.Data, error) {
 	pageData, err := getMultiPageData(html)
 	if err != nil {
-		return downloader.EmptyList, err
+		return nil, err
 	}
 	if !config.Playlist {
 		// handle URL that has a playlist, mainly for unified titles
@@ -189,6 +190,10 @@ func extractNormalVideo(url, html string) ([]downloader.Data, error) {
 			p, _ = strconv.Atoi(pageString[1])
 		}
 
+		if len(pageData.VideoData.Pages) < p || p < 1 {
+			return nil, extractors.ErrURLParseFailed
+		}
+
 		page := pageData.VideoData.Pages[p-1]
 		options := bilibiliOptions{
 			url:  url,
@@ -240,7 +245,7 @@ func Extract(url string) ([]downloader.Data, error) {
 	var err error
 	html, err := request.Get(url, referer, nil)
 	if err != nil {
-		return downloader.EmptyList, err
+		return nil, err
 	}
 	if strings.Contains(url, "bangumi") {
 		// handle bangumi
@@ -330,10 +335,13 @@ func bilibiliDownload(options bilibiliOptions) downloader.Data {
 		title = tempTitle
 	}
 
-	downloader.Caption(
+	err = downloader.Caption(
 		fmt.Sprintf("https://comment.bilibili.com/%s.xml", options.cid),
 		options.url, title, "xml",
 	)
+	if err != nil {
+		return downloader.EmptyData(options.url, err)
+	}
 
 	return downloader.Data{
 		Site:    "哔哩哔哩 bilibili.com",

extractors/defs.go

@@ -0,0 +1,7 @@
+package extractors
+
+import (
+	"errors"
+)
+
+var ErrURLParseFailed = errors.New("url parse failed")

extractors/douyin/douyin.go

@@ -2,6 +2,7 @@ package douyin
 
 import (
 	"github.com/iawia002/annie/downloader"
+	"github.com/iawia002/annie/extractors"
 	"github.com/iawia002/annie/request"
 	"github.com/iawia002/annie/utils"
 )
@@ -11,7 +12,7 @@ func Extract(url string) ([]downloader.Data, error) {
 	var err error
 	html, err := request.Get(url, url, nil)
 	if err != nil {
-		return downloader.EmptyList, err
+		return nil, err
 	}
 	var title string
 	desc := utils.MatchOneOf(html, `<p class="desc">(.+?)</p>`)
@@ -20,10 +21,15 @@ func Extract(url string) ([]downloader.Data, error) {
 	} else {
 		title = "抖音短视频"
 	}
-	realURL := utils.MatchOneOf(html, `playAddr: "(.+?)"`)[1]
+	realURLs := utils.MatchOneOf(html, `playAddr: "(.+?)"`)
+	if realURLs == nil || len(realURLs) < 2 {
+		return nil, extractors.ErrURLParseFailed
+	}
+	realURL := realURLs[1]
+
 	size, err := request.Size(realURL, url)
 	if err != nil {
-		return downloader.EmptyList, err
+		return nil, err
 	}
 	urlData := downloader.URL{
 		URL:  realURL,

extractors/douyu/douyu.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 
 	"github.com/iawia002/annie/downloader"
+	"github.com/iawia002/annie/extractors"
 	"github.com/iawia002/annie/request"
 	"github.com/iawia002/annie/utils"
 )
@@ -52,26 +53,35 @@ func Extract(url string) ([]downloader.Data, error) {
 	var err error
 	liveVid := utils.MatchOneOf(url, `https?://www.douyu.com/(\S+)`)
 	if liveVid != nil {
-		return downloader.EmptyList, errors.New("暂不支持斗鱼直播")
+		return nil, errors.New("暂不支持斗鱼直播")
 	}
 
 	html, err := request.Get(url, url, nil)
 	if err != nil {
-		return downloader.EmptyList, err
+		return nil, err
 	}
-	title := utils.MatchOneOf(html, `<title>(.*?)</title>`)[1]
+	titles := utils.MatchOneOf(html, `<title>(.*?)</title>`)
+	if titles == nil || len(titles) < 2 {
+		return nil, extractors.ErrURLParseFailed
+	}
+	title := titles[1]
+
+	vids := utils.MatchOneOf(url, `https?://v.douyu.com/show/(\S+)`)
+	if vids == nil || len(vids) < 2 {
+		return nil, extractors.ErrURLParseFailed
+	}
+	vid := vids[1]
 
-	vid := utils.MatchOneOf(url, `https?://v.douyu.com/show/(\S+)`)[1]
 	dataString, err := request.Get("http://vmobile.douyu.com/video/getInfo?vid="+vid, url, nil)
 	if err != nil {
-		return downloader.EmptyList, err
+		return nil, err
 	}
 	var dataDict douyuData
 	json.Unmarshal([]byte(dataString), &dataDict)
 
 	m3u8URLs, totalSize, err := douyuM3u8(dataDict.Data.VideoURL)
 	if err != nil {
-		return downloader.EmptyList, err
+		return nil, err
 	}
 	urls := make([]downloader.URL, len(m3u8URLs))
 	for index, u := range m3u8URLs {