You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
star/mget/util.go

142 lines
3.1 KiB
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
}