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