Skip to content

Commit 0c17093

Browse files
monkeyWiemonkeyWie
authored
feat: task separate proxy (#727)
1 parent 6538b29 commit 0c17093

File tree

22 files changed

+447
-67
lines changed

22 files changed

+447
-67
lines changed

‎internal/controller/controller.go‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,15 @@ package controller
22

33
import (
44
"github.com/GopeedLab/gopeed/pkg/base"
5+
"net/http"
6+
"net/url"
57
"os"
68
"path/filepath"
79
)
810

911
type Controller struct {
10-
GetConfig func(v any)
11-
ProxyConfig *base.DownloaderProxyConfig
12+
GetConfig func(v any)
13+
GetProxy func(requestProxy *base.RequestProxy) func(*http.Request) (*url.URL, error)
1214
FileController
1315
//ContextDialer() (proxy.Dialer, error)
1416
}
@@ -23,6 +25,7 @@ type DefaultFileController struct {
2325
func NewController() *Controller {
2426
return &Controller{
2527
GetConfig: func(v any) {},
28+
GetProxy: func(requestProxy *base.RequestProxy) func(*http.Request) (*url.URL, error) { return nil },
2629
FileController: &DefaultFileController{},
2730
}
2831
}

‎internal/protocol/bt/fetcher.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ func (f *Fetcher) initClient() (err error) {
7373
cfg.Bep20 = fmt.Sprintf("-GP%s-", parseBep20())
7474
cfg.ExtendedHandshakeClientVersion = fmt.Sprintf("Gopeed %s", base.Version)
7575
cfg.ListenPort = f.config.ListenPort
76-
cfg.HTTPProxy = f.ctl.ProxyConfig.ToHandler()
76+
cfg.HTTPProxy = f.ctl.GetProxy(f.meta.Req.Proxy)
7777
cfg.DefaultStorage = newFileOpts(newFileClientOpts{
7878
ClientBaseDir: cfg.DataDir,
7979
HandleFileTorrent: func(infoHash metainfo.Hash, ft *fileTorrentImpl) {
@@ -98,11 +98,14 @@ func (f *Fetcher) initClient() (err error) {
9898
}
9999

100100
func (f *Fetcher) Resolve(req *base.Request) error {
101+
if err := base.ParseReqExtra[bt.ReqExtra](req); err != nil {
102+
return err
103+
}
104+
f.meta.Req = req
101105
if err := f.addTorrent(req, false); err != nil {
102106
return err
103107
}
104108
f.updateRes()
105-
f.meta.Req = req
106109
return nil
107110
}
108111

@@ -339,9 +342,6 @@ func (f *Fetcher) WaitUpload() (err error) {
339342
}
340343

341344
func (f *Fetcher) addTorrent(req *base.Request, fromUpload bool) (err error) {
342-
if err = base.ParseReqExtra[bt.ReqExtra](req); err != nil {
343-
return
344-
}
345345
if err = f.initClient(); err != nil {
346346
return
347347
}

‎internal/protocol/bt/fetcher_test.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ import (
88
"github.com/GopeedLab/gopeed/internal/test"
99
"github.com/GopeedLab/gopeed/pkg/base"
1010
"github.com/GopeedLab/gopeed/pkg/protocol/bt"
11+
gohttp "net/http"
12+
"net/url"
1113
"os"
1214
"reflect"
1315
"testing"
@@ -210,7 +212,9 @@ func buildConfigFetcher(proxyConfig *base.DownloaderProxyConfig) fetcher.Fetcher
210212
newController.GetConfig = func(v any) {
211213
json.Unmarshal([]byte(test.ToJson(mockCfg)), v)
212214
}
213-
newController.ProxyConfig = proxyConfig
215+
newController.GetProxy = func(requestProxy *base.RequestProxy) func(*gohttp.Request) (*url.URL, error) {
216+
return proxyConfig.ToHandler()
217+
}
214218
fetcher.Setup(newController)
215219
return fetcher
216220
}

‎internal/protocol/http/fetcher.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ func (f *Fetcher) Resolve(req *base.Request) error {
7979
if err := base.ParseReqExtra[fhttp.ReqExtra](req); err != nil {
8080
return err
8181
}
82+
f.meta.Req = req
8283
httpReq, err := f.buildRequest(nil, req)
8384
if err != nil {
8485
return err
@@ -153,7 +154,6 @@ func (f *Fetcher) Resolve(req *base.Request) error {
153154
file.Name = httpReq.URL.Hostname()
154155
}
155156
res.Files = append(res.Files, file)
156-
f.meta.Req = req
157157
f.meta.Res = res
158158
return nil
159159
}
@@ -444,7 +444,7 @@ func (f *Fetcher) splitChunk() (chunks []*chunk) {
444444

445445
func (f *Fetcher) buildClient() *http.Client {
446446
transport := &http.Transport{
447-
Proxy: f.ctl.ProxyConfig.ToHandler(),
447+
Proxy: f.ctl.GetProxy(f.meta.Req.Proxy),
448448
}
449449
// Cookie handle
450450
jar, _ := cookiejar.New(nil)

‎internal/protocol/http/fetcher_test.go‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ import (
99
"github.com/GopeedLab/gopeed/pkg/base"
1010
"github.com/GopeedLab/gopeed/pkg/protocol/http"
1111
"net"
12+
gohttp "net/http"
13+
"net/url"
1214
"testing"
1315
"time"
1416
)
@@ -385,10 +387,12 @@ func downloadResume(listener net.Listener, connections int, t *testing.T) {
385387
func downloadWithProxy(httpListener net.Listener, proxyListener net.Listener, t *testing.T) {
386388
fetcher := downloadReady(httpListener, 4, t)
387389
ctl := controller.NewController()
388-
ctl.ProxyConfig = &base.DownloaderProxyConfig{
389-
Enable: true,
390-
Scheme: "socks5",
391-
Host: proxyListener.Addr().String(),
390+
ctl.GetProxy = func(requestProxy *base.RequestProxy) func(*gohttp.Request) (*url.URL, error) {
391+
return (&base.DownloaderProxyConfig{
392+
Enable: true,
393+
Scheme: "socks5",
394+
Host: proxyListener.Addr().String(),
395+
}).ToHandler()
392396
}
393397
fetcher.Setup(ctl)
394398
err := fetcher.Start()

‎pkg/base/model.go‎

Lines changed: 34 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ type Request struct {
1616
Extra any `json:"extra"`
1717
// Labels is used to mark the download task
1818
Labels map[string]string `json:"labels"`
19+
// Proxy is special proxy config for request
20+
Proxy *RequestProxy `json:"proxy"`
1921
}
2022

2123
func (r *Request) Validate() error {
@@ -25,6 +27,37 @@ func (r *Request) Validate() error {
2527
return nil
2628
}
2729

30+
type RequestProxyMode string
31+
32+
const (
33+
// RequestProxyModeFollow follow setting proxy
34+
RequestProxyModeFollow RequestProxyMode = "follow"
35+
// RequestProxyModeNone not use proxy
36+
RequestProxyModeNone RequestProxyMode = "none"
37+
// RequestProxyModeCustom custom proxy
38+
RequestProxyModeCustom RequestProxyMode = "custom"
39+
)
40+
41+
type RequestProxy struct {
42+
Mode RequestProxyMode `json:"mode"`
43+
Scheme string `json:"scheme"`
44+
Host string `json:"host"`
45+
Usr string `json:"usr"`
46+
Pwd string `json:"pwd"`
47+
}
48+
49+
func (p *RequestProxy) ToHandler() func(r *http.Request) (*url.URL, error) {
50+
if p == nil || p.Mode != RequestProxyModeCustom {
51+
return nil
52+
}
53+
54+
if p.Scheme == "" || p.Host == "" {
55+
return nil
56+
}
57+
58+
return http.ProxyURL(util.BuildProxyUrl(p.Scheme, p.Host, p.Usr, p.Pwd))
59+
}
60+
2861
// Resource download resource
2962
type Resource struct {
3063
// if name is not empty, the resource is a folder and the name is the folder name
@@ -94,15 +127,7 @@ func (o *Options) InitSelectFiles(fileSize int) {
94127
}
95128

96129
func (o *Options) Clone() *Options {
97-
if o == nil {
98-
return nil
99-
}
100-
return &Options{
101-
Name: o.Name,
102-
Path: o.Path,
103-
SelectFiles: o.SelectFiles,
104-
Extra: o.Extra,
105-
}
130+
return util.DeepClone(o)
106131
}
107132

108133
func ParseReqExtra[E any](req *Request) error {

‎pkg/download/downloader.go‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ import (
1212
"github.com/rs/zerolog"
1313
"github.com/rs/zerolog/pkgerrors"
1414
"github.com/virtuald/go-paniclog"
15+
gohttp "net/http"
16+
"net/url"
1517
"os"
1618
"path/filepath"
1719
"sort"
@@ -253,7 +255,20 @@ func (d *Downloader) setupFetcher(fm fetcher.FetcherManager, fetcher fetcher.Fet
253255
ctl.GetConfig = func(v any) {
254256
d.getProtocolConfig(fm.Name(), v)
255257
}
256-
ctl.ProxyConfig = d.cfg.Proxy
258+
// Get proxy config, task request proxy config has higher priority, then use global proxy config
259+
ctl.GetProxy = func(requestProxy *base.RequestProxy) func(*gohttp.Request) (*url.URL, error) {
260+
if requestProxy == nil {
261+
return d.cfg.Proxy.ToHandler()
262+
}
263+
switch requestProxy.Mode {
264+
case base.RequestProxyModeNone:
265+
return nil
266+
case base.RequestProxyModeCustom:
267+
return requestProxy.ToHandler()
268+
default:
269+
return d.cfg.Proxy.ToHandler()
270+
}
271+
}
257272
fetcher.Setup(ctl)
258273
}
259274

‎pkg/download/downloader_test.go‎

Lines changed: 41 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -159,21 +159,21 @@ func TestDownloader_CreateDirectBatch(t *testing.T) {
159159

160160
func TestDownloader_CreateWithProxy(t *testing.T) {
161161
// No proxy
162-
doTestDownloaderCreateWithProxy(t, false, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
162+
doTestDownloaderCreateWithProxy(t, false, nil, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
163163
return nil
164164
}, nil)
165165
// Disable proxy
166-
doTestDownloaderCreateWithProxy(t, false, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
166+
doTestDownloaderCreateWithProxy(t, false, nil, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
167167
proxyCfg.Enable = false
168168
return proxyCfg
169169
}, nil)
170170
// Enable system proxy but not set proxy environment variable
171-
doTestDownloaderCreateWithProxy(t, false, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
171+
doTestDownloaderCreateWithProxy(t, false, nil, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
172172
proxyCfg.System = true
173173
return proxyCfg
174174
}, nil)
175175
// Enable proxy but error proxy environment variable
176-
doTestDownloaderCreateWithProxy(t, false, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
176+
doTestDownloaderCreateWithProxy(t, false, nil, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
177177
os.Setenv("HTTP_PROXY", "http://127.0.0.1:1234")
178178
os.Setenv("HTTPS_PROXY", "http://127.0.0.1:1234")
179179
proxyCfg.System = true
@@ -184,33 +184,50 @@ func TestDownloader_CreateWithProxy(t *testing.T) {
184184
}
185185
})
186186
// Enable system proxy and set proxy environment variable
187-
doTestDownloaderCreateWithProxy(t, false, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
187+
doTestDownloaderCreateWithProxy(t, false, nil, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
188188
os.Setenv("HTTP_PROXY", proxyCfg.ToUrl().String())
189189
os.Setenv("HTTPS_PROXY", proxyCfg.ToUrl().String())
190190
proxyCfg.System = true
191191
return proxyCfg
192192
}, nil)
193193
// Invalid proxy scheme
194-
doTestDownloaderCreateWithProxy(t, false, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
194+
doTestDownloaderCreateWithProxy(t, false, nil, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
195195
proxyCfg.Scheme = ""
196196
return proxyCfg
197197
}, nil)
198198
// Invalid proxy host
199-
doTestDownloaderCreateWithProxy(t, false, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
199+
doTestDownloaderCreateWithProxy(t, false, nil, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
200200
proxyCfg.Host = ""
201201
return proxyCfg
202202
}, nil)
203203
// Use proxy without auth
204-
doTestDownloaderCreateWithProxy(t, false, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
204+
doTestDownloaderCreateWithProxy(t, false, nil, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
205205
return proxyCfg
206206
}, nil)
207207
// Use proxy with auth
208-
doTestDownloaderCreateWithProxy(t, true, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
208+
doTestDownloaderCreateWithProxy(t, true, nil, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig {
209209
return proxyCfg
210210
}, nil)
211+
212+
// Request proxy mode follow
213+
doTestDownloaderCreateWithProxy(t, false, func(reqProxy *base.RequestProxy) *base.RequestProxy {
214+
reqProxy.Mode = base.RequestProxyModeFollow
215+
return reqProxy
216+
}, nil, nil)
217+
218+
// Request proxy mode none
219+
doTestDownloaderCreateWithProxy(t, false, func(reqProxy *base.RequestProxy) *base.RequestProxy {
220+
reqProxy.Mode = base.RequestProxyModeNone
221+
return reqProxy
222+
}, nil, nil)
223+
224+
// Request proxy mode custom
225+
doTestDownloaderCreateWithProxy(t, false, func(reqProxy *base.RequestProxy) *base.RequestProxy {
226+
return reqProxy
227+
}, nil, nil)
211228
}
212229

213-
func doTestDownloaderCreateWithProxy(t *testing.T, auth bool, buildProxyConfig func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig, errHandler func(err error)) {
230+
func doTestDownloaderCreateWithProxy(t *testing.T, auth bool, buildReqProxy func(reqProxy *base.RequestProxy) *base.RequestProxy, buildProxyConfig func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig, errHandler func(err error)) {
214231
usr, pwd := "", ""
215232
if auth {
216233
usr, pwd = "admin", "123"
@@ -223,17 +240,29 @@ func doTestDownloaderCreateWithProxy(t *testing.T, auth bool, buildProxyConfig f
223240
t.Fatal(err)
224241
}
225242
defer downloader.Clear()
226-
downloader.cfg.DownloaderStoreConfig.Proxy = buildProxyConfig(&base.DownloaderProxyConfig{
243+
globalProxyCfg := &base.DownloaderProxyConfig{
227244
Enable: true,
228245
Scheme: "socks5",
229246
Host: proxyListener.Addr().String(),
230247
Usr: usr,
231248
Pwd: pwd,
232-
})
249+
}
250+
if buildProxyConfig != nil {
251+
globalProxyCfg = buildProxyConfig(globalProxyCfg)
252+
}
253+
downloader.cfg.DownloaderStoreConfig.Proxy = globalProxyCfg
233254

234255
req := &base.Request{
235256
URL: test.ExternalDownloadUrl,
236257
}
258+
if buildReqProxy != nil {
259+
req.Proxy = buildReqProxy(&base.RequestProxy{
260+
Scheme: "socks5",
261+
Host: proxyListener.Addr().String(),
262+
Usr: usr,
263+
Pwd: pwd,
264+
})
265+
}
237266
rr, err := downloader.Resolve(req)
238267
if err != nil {
239268
if errHandler == nil {

‎pkg/download/engine/engine.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ func (e *Engine) await(value any) {
106106
}
107107

108108
func (e *Engine) Close() {
109-
e.loop.Stop()
109+
e.loop.StopNoWait()
110110
}
111111

112112
type Config struct {

‎pkg/download/extension.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -694,8 +694,12 @@ type ExtensionTask struct {
694694
}
695695

696696
func NewExtensionTask(download *Downloader, task *Task) *ExtensionTask {
697+
newTask := task.clone()
698+
// Assign the pointer of the properties that the extension supports modification
699+
newTask.Meta = task.Meta
700+
newTask.Status = task.Status
697701
return &ExtensionTask{
698-
Task: task.clone(),
702+
Task: newTask,
699703
download: download,
700704
}
701705
}

0 commit comments

Comments
 (0)