support to download as stream

LinuxSuRen · LinuxSuRen · commit 2f36367014e9 · 2025-04-21T08:46:20.000+08:00
diff --git a/pkg/net/http.go b/pkg/net/http.go
@@ -97,8 +97,7 @@ func (h *HTTPDownloader) fetchProxyFromEnv(scheme string) {
 	}
 }
 
-// DownloadFile download a file with the progress
-func (h *HTTPDownloader) DownloadFile() error {
+func (h *HTTPDownloader) DownloadAsStream(writer io.Writer) (err error) {
 	filepath, downloadURL, showProgress := h.TargetFilePath, h.URL, h.ShowProgress
 	// Get the data
 	if h.Context == nil {
@@ -181,22 +180,32 @@ func (h *HTTPDownloader) DownloadFile() error {
 		}
 	}
 
-	if err := os.MkdirAll(path.Dir(filepath), os.FileMode(0755)); err != nil {
-		return err
+	h.progressIndicator.Writer = writer
+	h.progressIndicator.Init()
+
+	// Write the body to file
+	_, err = io.Copy(h.progressIndicator, resp.Body)
+	return
+}
+
+// DownloadFile download a file with the progress
+func (h *HTTPDownloader) DownloadFile() (err error) {
+	filepath := h.TargetFilePath
+	if err = os.MkdirAll(path.Dir(filepath), os.FileMode(0755)); err != nil {
+		return
 	}
 
 	// Create the file
-	out, err := os.Create(filepath)
+	var out io.WriteCloser
+	out, err = os.Create(filepath)
 	if err != nil {
-		_ = out.Close()
-		return err
+		return
 	}
+	defer func() {
+		_ = out.Close()
+	}()
 
-	h.progressIndicator.Writer = out
-	h.progressIndicator.Init()
-
-	// Write the body to file
-	_, err = io.Copy(h.progressIndicator, resp.Body)
+	err = h.DownloadAsStream(out)
 	return err
 }
 
@@ -272,6 +281,38 @@ func (c *ContinueDownloader) WithBasicAuth(username, password string) *ContinueD
 	return c
 }
 
+func (c *ContinueDownloader) DownloadWithContinueAsStream(targetURL string, output io.Writer, index, continueAt, end int64, showProgress bool) (err error) {
+	c.downloader = &HTTPDownloader{
+		URL:                targetURL,
+		ShowProgress:       showProgress,
+		NoProxy:            c.noProxy,
+		RoundTripper:       c.roundTripper,
+		InsecureSkipVerify: c.insecureSkipVerify,
+		UserName:           c.UserName,
+		Password:           c.Password,
+		Context:            c.Context,
+		Timeout:            c.Timeout,
+	}
+	if index >= 0 {
+		c.downloader.Title = fmt.Sprintf("Downloading part %d", index)
+	}
+
+	if continueAt >= 0 {
+		c.downloader.Header = make(map[string]string, 1)
+
+		if end > continueAt {
+			c.downloader.Header["Range"] = fmt.Sprintf("bytes=%d-%d", continueAt, end)
+		} else {
+			c.downloader.Header["Range"] = fmt.Sprintf("bytes=%d-", continueAt)
+		}
+	}
+
+	if err = c.downloader.DownloadAsStream(output); err != nil {
+		err = fmt.Errorf("cannot download from %s, error: %v", targetURL, err)
+	}
+	return
+}
+
 // DownloadWithContinue downloads the files continuously
 func (c *ContinueDownloader) DownloadWithContinue(targetURL, output string, index, continueAt, end int64, showProgress bool) (err error) {
 	c.downloader = &HTTPDownloader{
@@ -306,6 +347,44 @@ func (c *ContinueDownloader) DownloadWithContinue(targetURL, output string, inde
 	return
 }
 
+func DetectSizeWithRoundTripperAndAuthStream(targetURL string, output io.Writer, showProgress, noProxy, insecureSkipVerify bool,
+	roundTripper http.RoundTripper, username, password string, timeout time.Duration) (total int64, rangeSupport bool, err error) {
+	downloader := HTTPDownloader{
+		URL:                targetURL,
+		ShowProgress:       showProgress,
+		RoundTripper:       roundTripper,
+		NoProxy:            false, // below HTTP request does not need proxy
+		InsecureSkipVerify: insecureSkipVerify,
+		UserName:           username,
+		Password:           password,
+		Timeout:            timeout,
+	}
+
+	var detectOffset int64
+	var lenErr error
+
+	detectOffset = 2
+	downloader.Header = make(map[string]string, 1)
+	downloader.Header["Range"] = fmt.Sprintf("bytes=%d-", detectOffset)
+
+	downloader.PreStart = func(resp *http.Response) bool {
+		rangeSupport = resp.StatusCode == http.StatusPartialContent
+		contentLen := resp.Header.Get("Content-Length")
+		if total, lenErr = strconv.ParseInt(contentLen, 10, 0); lenErr == nil {
+			total += detectOffset
+		} else {
+			rangeSupport = false
+		}
+		//  always return false because we just want to get the header from response
+		return false
+	}
+
+	if err = downloader.DownloadAsStream(output); err != nil || lenErr != nil {
+		err = fmt.Errorf("cannot download from %s, response error: %v, content length error: %v", targetURL, err, lenErr)
+	}
+	return
+}
+
 // DetectSizeWithRoundTripperAndAuth returns the size of target resource
 func DetectSizeWithRoundTripperAndAuth(targetURL, output string, showProgress, noProxy, insecureSkipVerify bool,
 	roundTripper http.RoundTripper, username, password string, timeout time.Duration) (total int64, rangeSupport bool, err error) {
diff --git a/pkg/net/multi_thread.go b/pkg/net/multi_thread.go
@@ -3,6 +3,7 @@ package net
 import (
 	"context"
 	"fmt"
+	"io"
 	"net/http"
 	"os"
 	"os/signal"
@@ -75,6 +76,120 @@ func (d *MultiThreadDownloader) WithBearerToken(bearerToken string) *MultiThread
 	return d
 }
 
+func (d *MultiThreadDownloader) DownloadWithContext(ctx context.Context, targetURL string, outputWriter io.Writer, thread int) (err error) {
+	// get the total size of the target file
+	var total int64
+	var rangeSupport bool
+	if total, rangeSupport, err = DetectSizeWithRoundTripperAndAuthStream(targetURL, outputWriter, d.showProgress,
+		d.noProxy, d.insecureSkipVerify, d.roundTripper, d.username, d.password, d.timeout); rangeSupport && err != nil {
+		return
+	}
+
+	if rangeSupport {
+		unit := total / int64(thread)
+		offset := total - unit*int64(thread)
+		var wg sync.WaitGroup
+		var partItems map[int]string
+		var m sync.Mutex
+
+		defer func() {
+			// remove all partial files
+			for _, part := range partItems {
+				_ = os.RemoveAll(part)
+			}
+		}()
+
+		c := make(chan os.Signal, 1)
+		signal.Notify(c, os.Interrupt)
+		ctx, cancel := context.WithCancel(context.Background())
+		var canceled bool
+
+		go func() {
+			<-c
+			canceled = true
+			cancel()
+		}()
+
+		fmt.Printf("start to download with %d threads, size: %d, unit: %d", thread, total, unit)
+		for i := 0; i < thread; i++ {
+			fmt.Println() // TODO take position, should take over by progerss bars
+			wg.Add(1)
+			go func(index int, wg *sync.WaitGroup, ctx context.Context) {
+				defer wg.Done()
+				outputFile, err := os.CreateTemp(os.TempDir(), fmt.Sprintf("part-%d", index))
+				if err != nil {
+					fmt.Println("failed to create template file", err)
+				}
+				outputFile.Close()
+
+				m.Lock()
+				partItems[index] = outputFile.Name()
+				m.Unlock()
+
+				end := unit*int64(index+1) - 1
+				if index == thread-1 {
+					// this is the last part
+					end += offset
+				}
+				start := unit * int64(index)
+
+				downloader := &ContinueDownloader{}
+				downloader.WithoutProxy(d.noProxy).
+					WithRoundTripper(d.roundTripper).
+					WithInsecureSkipVerify(d.insecureSkipVerify).
+					WithBasicAuth(d.username, d.password).
+					WithContext(ctx).WithTimeout(d.timeout)
+				if downloadErr := downloader.DownloadWithContinue(targetURL, outputFile.Name(),
+					int64(index), start, end, d.showProgress); downloadErr != nil {
+					fmt.Println(downloadErr)
+				}
+			}(i, &wg, ctx)
+		}
+
+		wg.Wait()
+		// ProgressIndicator{}.Close()
+		if canceled {
+			err = fmt.Errorf("download process canceled")
+			return
+		}
+
+		// make the cursor right
+		// TODO the progress component should take over it
+		if thread > 1 {
+			// line := GetCurrentLine()
+			time.Sleep(time.Second)
+			fmt.Printf("\033[%dE\n", thread) // move to the target line
+			time.Sleep(time.Second * 5)
+		}
+
+		for i := 0; i < thread; i++ {
+			partFile := partItems[i]
+			if data, ferr := os.ReadFile(partFile); ferr == nil {
+				if _, err = outputWriter.Write(data); err != nil {
+					err = fmt.Errorf("failed to write file: '%s'", partFile)
+					break
+				} else if !d.keepParts {
+					_ = os.RemoveAll(partFile)
+				}
+			} else {
+				err = fmt.Errorf("failed to read file: '%s'", partFile)
+				break
+			}
+		}
+	} else {
+		fmt.Println("cannot download it using multiple threads, failed to one")
+		downloader := &ContinueDownloader{}
+		downloader.WithoutProxy(d.noProxy)
+		downloader.WithRoundTripper(d.roundTripper)
+		downloader.WithInsecureSkipVerify(d.insecureSkipVerify)
+		downloader.WithTimeout(d.timeout)
+		downloader.WithBasicAuth(d.username, d.password)
+		err = downloader.DownloadWithContinueAsStream(targetURL, outputWriter, -1, 0, 0, true)
+		d.suggestedFilename = downloader.GetSuggestedFilename()
+	}
+	return
+}
+
 // Download starts to download the target URL
 func (d *MultiThreadDownloader) Download(targetURL, targetFilePath string, thread int) (err error) {
 	// get the total size of the target file
