star/mget/util.go

package mget

import (
	"b612.me/staros"
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"regexp"
	"runtime"
	"sync/atomic"
	"time"
)

func parseContentRange(contentRange string) (start, end, total int64, err error) {
	_, err = fmt.Sscanf(contentRange, "bytes %d-%d/%d", &start, &end, &total)
	return
}

func GetFileName(resp *http.Response) string {
	fname := getFileName(resp)
	var idx = 0
	for {
		idx++
		if staros.Exists(fname) {
			if staros.Exists(fname + ".bgrd") {
				return fname
			}
			fname = fmt.Sprintf("%s.%d", fname, idx)
		} else {
			break
		}
	}
	return fname
}

func getFileName(resp *http.Response) string {
	// 尝试从Content-Disposition头中提取文件名
	contentDisposition := resp.Header.Get("Content-Disposition")
	if contentDisposition != "" {
		// 使用正则表达式提取文件名
		re := regexp.MustCompile(`(?i)^attachment; filename="?(?P<filename>[^;"]+)`)
		matches := re.FindStringSubmatch(contentDisposition)
		if len(matches) > 1 {
			// 提取命名的捕获组
			for i, name := range re.SubexpNames() {
				if name == "filename" {
					return matches[i]
				}
			}
		}
	}
	// 提取路径中的最后一个元素作为文件名
	return path.Base(resp.Request.URL.Path)
}

func IOWriter(stopCtx context.Context, ch chan Buffer, state *uint32, di *downloadinfo, reader io.ReadCloser, bufSize int, start *int64, end *int64) error {
	defer reader.Close()
	for {
		buf := make([]byte, bufSize)
		select {
		case <-stopCtx.Done():
			return nil
		default:
			if atomic.LoadUint32(state) == 1 {
				runtime.Gosched()
				time.Sleep(time.Millisecond)
				continue
			}
			n, err := reader.Read(buf)
			if n > 0 {
				ch <- Buffer{Data: buf[:n], Start: uint64(*start)}
				*start += int64(n)
				di.AddCurrent(int64(n))
			}
			if *end != 0 && *start >= *end {
				return nil
			}
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return err
			}
		}
	}
}

func createFileWithSize(filename string, size int64) (*os.File, error) {
	file, err := os.Create(filename)
	if err != nil {
		return nil, err
	}
	if size == 0 {
		return file, nil
	}
	// 调整文件指针到指定大小位置
	if _, err = file.Seek(size-1, 0); err != nil {
		return nil, err
	}

	// 写入一个空字节，以确保文件达到所需大小
	if _, err = file.Write([]byte{0}); err != nil {
		return nil, err
	}

	return file, nil
}

func CloneHeader(original http.Header) http.Header {
	newHeader := make(http.Header)
	for key, values := range original {
		copiedValues := make([]string, len(values))
		copy(copiedValues, values)
		newHeader[key] = copiedValues
	}
	return newHeader
}

func CloneCookies(original []*http.Cookie) []*http.Cookie {
	cloned := make([]*http.Cookie, len(original))
	for i, cookie := range original {
		cloned[i] = &http.Cookie{
			Name:       cookie.Name,
			Value:      cookie.Value,
			Path:       cookie.Path,
			Domain:     cookie.Domain,
			Expires:    cookie.Expires,
			RawExpires: cookie.RawExpires,
			MaxAge:     cookie.MaxAge,
			Secure:     cookie.Secure,
			HttpOnly:   cookie.HttpOnly,
			SameSite:   cookie.SameSite,
			Raw:        cookie.Raw,
			Unparsed:   append([]string(nil), cookie.Unparsed...),
		}
	}
	return cloned
}