feat: Added header cache, optional header settings...

This commit is contained in:
putyy
2025-04-23 16:55:24 +08:00
parent 0bb1a21a76
commit 6086bd7086
12 changed files with 92 additions and 23 deletions

View File

@@ -24,6 +24,7 @@ type Config struct {
WxAction bool `json:"WxAction"`
TaskNumber int `json:"TaskNumber"`
UserAgent string `json:"UserAgent"`
UseHeaders string `json:"UseHeaders"`
}
func initConfig() *Config {
@@ -43,7 +44,8 @@ func initConfig() *Config {
"AutoProxy": true,
"WxAction": true,
"TaskNumber": __TaskNumber__,
"UserAgent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36"
"UserAgent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36",
"UseHeaders": "User-Agent,Referer,Authorization,Cookie"
}
`
def = strings.ReplaceAll(def, "__TaskNumber__", strconv.Itoa(runtime.NumCPU()*2))
@@ -51,9 +53,23 @@ func initConfig() *Config {
storage: NewStorage("config.json", []byte(def)),
}
defaultMap := make(map[string]interface{})
_ = json.Unmarshal([]byte(def), &defaultMap)
data, err := globalConfig.storage.Load()
if err == nil {
_ = json.Unmarshal(data, &globalConfig)
var loadedMap map[string]interface{}
_ = json.Unmarshal(data, &loadedMap)
for key, val := range defaultMap {
if _, ok := loadedMap[key]; !ok {
loadedMap[key] = val
}
}
finalBytes, _ := json.Marshal(loadedMap)
_ = json.Unmarshal(finalBytes, &globalConfig)
} else {
globalLogger.Esg(err, "load config err")
}
@@ -77,6 +93,7 @@ func (c *Config) setConfig(config Config) {
c.AutoProxy = config.AutoProxy
c.TaskNumber = config.TaskNumber
c.WxAction = config.WxAction
c.UseHeaders = config.UseHeaders
if oldProxy != c.UpstreamProxy {
proxyOnce.setTransport()
}

View File

@@ -5,7 +5,6 @@ import (
"io"
"log"
"net/http"
"net/http/cookiejar"
"net/url"
"os"
"path/filepath"
@@ -32,17 +31,19 @@ type FileDownloader struct {
totalTasks int
TotalSize int64
IsMultiPart bool
Headers map[string]string
DownloadTaskList []*DownloadTask
progressCallback ProgressCallback
}
func NewFileDownloader(url, filename string, totalTasks int) *FileDownloader {
func NewFileDownloader(url, filename string, totalTasks int, headers map[string]string) *FileDownloader {
return &FileDownloader{
Url: url,
FileName: filename,
totalTasks: totalTasks,
IsMultiPart: false,
TotalSize: 0,
Headers: headers,
DownloadTaskList: make([]*DownloadTask, 0),
}
}
@@ -53,10 +54,16 @@ func (fd *FileDownloader) buildClient() *http.Client {
transport.Proxy = http.ProxyURL(fd.ProxyUrl)
}
// Cookie handle
jar, _ := cookiejar.New(nil)
return &http.Client{
Transport: transport,
Jar: jar,
}
}
func (fd *FileDownloader) setHeaders(request *http.Request) {
for key, values := range fd.Headers {
if strings.Contains(globalConfig.UseHeaders, key) {
request.Header.Set(key, values)
}
}
}
@@ -76,21 +83,22 @@ func (fd *FileDownloader) init() error {
}
}
req, err := http.NewRequest("HEAD", fd.Url, nil)
request, err := http.NewRequest("HEAD", fd.Url, nil)
if err != nil {
return fmt.Errorf("create request failed")
}
// 设置请求头
if globalConfig.UserAgent != "" {
req.Header.Set("User-Agent", globalConfig.UserAgent)
if _, ok := fd.Headers["User-Agent"]; !ok {
fd.Headers["User-Agent"] = globalConfig.UserAgent
}
if fd.Referer != "" {
req.Header.Set("Referer", fd.Referer)
if _, ok := fd.Headers["Referer"]; !ok {
fd.Headers["Referer"] = fd.Referer
}
resp, err := fd.buildClient().Do(req)
fd.setHeaders(request)
resp, err := fd.buildClient().Do(request)
if err != nil {
return fmt.Errorf("request failed" + err.Error())
}
@@ -182,8 +190,9 @@ func (fd *FileDownloader) startDownloadTask(waitGroup *sync.WaitGroup, progressC
globalLogger.Error().Stack().Err(err).Msgf("任务%d创建请求出错", task.taskID)
return
}
request.Header.Set("User-Agent", globalConfig.UserAgent)
request.Header.Set("Referer", fd.Referer)
fd.setHeaders(request)
if fd.IsMultiPart {
request.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", task.rangeStart, task.rangeEnd))
}

View File

@@ -336,6 +336,12 @@ func (p *Proxy) httpResponseEvent(resp *http.Response, ctx *goproxy.ProxyCtx) *h
Description: "",
ContentType: resp.Header.Get("Content-Type"),
}
// Store entire request headers as JSON
if headers, err := json.Marshal(resp.Request.Header); err == nil {
res.OtherData["headers"] = string(headers)
}
resourceOnce.markMedia(urlSign)
httpServerOnce.send("newResources", res)
}

View File

@@ -2,6 +2,8 @@ package core
import (
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/url"
"os"
@@ -142,7 +144,9 @@ func (r *Resource) download(mediaInfo MediaInfo, decodeStr string) {
}
}
downloader := NewFileDownloader(rawUrl, mediaInfo.SavePath, globalConfig.TaskNumber)
headers, _ := r.parseHeaders(mediaInfo)
downloader := NewFileDownloader(rawUrl, mediaInfo.SavePath, globalConfig.TaskNumber, headers)
downloader.progressCallback = func(totalDownloaded float64) {
r.progressEventsEmit(mediaInfo, strconv.Itoa(int(totalDownloaded))+"%", DownloadStatusRunning)
}
@@ -162,6 +166,27 @@ func (r *Resource) download(mediaInfo MediaInfo, decodeStr string) {
}(mediaInfo)
}
// 解析并组装 headers
func (r *Resource) parseHeaders(mediaInfo MediaInfo) (map[string]string, error) {
headers := make(map[string]string)
if hh, ok := mediaInfo.OtherData["headers"]; ok {
var tempHeaders map[string][]string
// 解析 JSON 字符串为 map[string][]string
if err := json.Unmarshal([]byte(hh), &tempHeaders); err != nil {
return headers, fmt.Errorf("parse headers JSON err: %v", err)
}
for key, values := range tempHeaders {
if len(values) > 0 {
headers[key] = values[0]
}
}
}
return headers, nil
}
func (r *Resource) wxFileDecode(mediaInfo MediaInfo, fileName, decodeStr string) (string, error) {
sourceFile, err := os.Open(fileName)
if err != nil {