diff --git a/agent/toolcatalog.go b/agent/toolcatalog.go index d564efe..3c018d9 100644 --- a/agent/toolcatalog.go +++ b/agent/toolcatalog.go @@ -36,7 +36,7 @@ func builtinToolsByAgent() map[string][]actool.CoreTool { return map[string][]actool.CoreTool{ "mainagent": ts.MainAgentTools(), "planner": ts.PlannerTools(), - "worker": ts.WorkerTools(), + "worker": ts.WorkerTools(), // 含 list_proxies（代理池只读，随开关生效） // goals（目标拆解器）默认绑 set_goals：它靠这个工具把拆出的目标写进图。 // 与 mainagent 共用同一受管工具，web 端可改描述/schema、按 agent 勾选。 "goals": {ts.setGoals()}, diff --git a/agent/tools.go b/agent/tools.go index 53073d0..4b57f79 100644 --- a/agent/tools.go +++ b/agent/tools.go @@ -53,6 +53,7 @@ func compactIntents(ns []*db.Node, parentsOf, yieldsOf map[int64][]int64) []map[ type ToolSet struct { as *db.AssetStore // asset store (optional; nil = asset tools not available) cs *db.CompanyStore // company store (optional) + ps *db.ProxyStore // proxy pool store (optional; nil = list_proxies unavailable) ts *db.ExplorationStore worker string taskID int64 // PG tasks.id; 0 when unknown (tests / orchestrator cross-task reads) diff --git a/agent/tools_insert.go b/agent/tools_insert.go index 9d0eca4..4c0593b 100644 --- a/agent/tools_insert.go +++ b/agent/tools_insert.go @@ -21,6 +21,9 @@ func (t *ToolSet) SetAssetStore(as *db.AssetStore, cs *db.CompanyStore) { t.cs = cs } +// SetProxyStore wires the proxy pool store so the list_proxies tool is active. +func (t *ToolSet) SetProxyStore(ps *db.ProxyStore) { t.ps = ps } + // assetInputItem is one element of the insert_assets "assets" array. type assetInputItem struct { Type string `json:"type"` // root_domain|ip|subdomain|app|service|endpoint @@ -160,7 +163,7 @@ func (t *ToolSet) insertAssets() actool.CoreTool { // 覆盖度相关：该资产是否与当前测试任务有关。 "related": map[string]any{ "type": "boolean", - "description": "该资产是否与【当前测试任务】相关：true(默认)才自动纳入资产覆盖度(测试范围分母)；false 则只入库、不计入本任务覆盖度(如顺带发现的旁站/无关资产)。仅在任务开启资产覆盖度功能时生效。", + "description": "该资产是否与【当前测试任务】相关：true(默认)才自动纳入资产覆盖度(测试范围分母)，该资产会进入待测资产中，如果是和任务无关的，例如CDN仅存储静态资源类，必须设置为false或不进行资产插入；false 则只入库、不计入本任务覆盖度(如顺带发现的旁站/无关资产)。仅在任务开启资产覆盖度功能时生效。", }, }, "type"), }, @@ -582,6 +585,66 @@ func (t *ToolSet) listCompanies() actool.CoreTool { ) } +// listProxies lets an agent enumerate healthy outbound proxies it can route through +// on its own (e.g. curl -x, proxychains) for a specific target. Read-only: it only +// lists nodes, it does not change any global egress setting. +func (t *ToolSet) listProxies() actool.CoreTool { + return readTool( + "list_proxies", + "列出代理池中【当前健康】的出口代理节点，供你在命令里自行使用（如 curl -x 、"+ + "proxychains、nmap --proxies）访问目标、轮换出口 IP。返回 http/https/socks5 代理的 "+ + "address（scheme://host:port）、地区、匿名度、延迟。可选 protocol/region/tag 过滤。"+ + "只读：不改变全局出口设置。代理池未开启时返回空。", + obj(map[string]any{ + "protocol": str("按协议过滤(可选)：http/https/socks5"), + "region": str("按地区码过滤(可选)，如 CN/US"), + "tag": str("按标签过滤(可选)"), + }), + func(_ context.Context, in json.RawMessage) (actool.Result, error) { + if t.ps == nil || !t.ps.PoolEnabled() { + return actool.Errorf("list_proxies 未启用: 代理池功能已关闭（在系统设置开启后可用）"), nil + } + var a struct { + Protocol string `json:"protocol"` + Region string `json:"region"` + Tag string `json:"tag"` + } + _ = json.Unmarshal(in, &a) + f := db.ProxyFilter{ + Protocol: strings.TrimSpace(a.Protocol), + Region: strings.TrimSpace(a.Region), + OnlyEnabled: true, + OnlyHealthy: true, + } + if tag := strings.TrimSpace(a.Tag); tag != "" { + f.Tags = []string{tag} + } + proxies, err := t.ps.ListProxies(f) + if err != nil { + return actool.Errorf("查询代理失败: " + err.Error()), nil + } + type proxyOut struct { + Address string `json:"address"` // scheme://host:port (含认证，供命令直接使用) + Protocol string `json:"protocol"` + Region string `json:"region,omitempty"` + Anonymity string `json:"anonymity,omitempty"` + LatencyMs int `json:"latency_ms"` + } + out := make([]proxyOut, 0, len(proxies)) + for _, p := range proxies { + out = append(out, proxyOut{ + Address: p.URL().String(), + Protocol: p.Protocol, + Region: p.Region, + Anonymity: p.Anonymity, + LatencyMs: p.LatencyMs, + }) + } + return jsonResult(map[string]any{"count": len(out), "proxies": out}) + }, + ) +} + // splitLines splits a multi-line string into non-empty trimmed lines. func splitLines(s string) []string { var out []string @@ -603,6 +666,8 @@ func (t *ToolSet) WorkerTools() []actool.CoreTool { t.searchAllWorkerTraces(), t.listWorkerTraces(), t.getWorkerTrace(), // asset management (handlers guard nil store internally) t.insertAssets(), t.addCompanyScope(), t.listAssets(), t.listCompanies(), + // outbound proxy pool (read-only; guarded when pool store/switch off) + t.listProxies(), } } diff --git a/agent/worker.go b/agent/worker.go index 99125c9..762443b 100644 --- a/agent/worker.go +++ b/agent/worker.go @@ -258,6 +258,7 @@ func (w *Worker) Execute(ctx context.Context, name string, taskID int64, as *db. tsx.SetCoverageEnabled(coverageEnabled) if as != nil { tsx.SetAssetStore(as, as.Companies()) + tsx.SetProxyStore(as.Proxies()) // list_proxies（只读，随代理池开关生效） } tsx.SetOwnerNode(intent.ID) // assets this worker discovers anchor to its intent → visible to the task tsx.SetEnrich(enr) // async DNS/HTTP auto-completion for assets this worker writes diff --git a/cmd/artex/main.go b/cmd/artex/main.go index 6023245..9ea0012 100644 --- a/cmd/artex/main.go +++ b/cmd/artex/main.go @@ -46,6 +46,7 @@ func main() { addr = flag.String("addr", ":8787", "HTTP listen address") dataDir = flag.String("data", filepath.Join(config.BaseDir(), "data"), "data directory for SQLite stores (default: data/ next to the executable)") proxy = flag.String("proxy", ":8788", "traffic recording proxy address (empty to disable)") + proxyGW = flag.String("proxy-pool-gateway", "127.0.0.1:8789", "proxy-pool forwarding gateway address (empty to disable 入口C)") ) flag.Parse() @@ -76,7 +77,7 @@ func main() { ctx, shutdown := shutdownContext(sigCtx) defer shutdown(agent.AbortShutdown) - mgr, err := server.NewManager(*dataDir, *proxy) + mgr, err := server.NewManager(*dataDir, *proxy, *proxyGW) if err != nil { log.Fatalf("open stores: %v", err) } diff --git a/db/assets.go b/db/assets.go index 61be527..a1fd248 100644 --- a/db/assets.go +++ b/db/assets.go @@ -94,6 +94,10 @@ func (d *DB) Assets() *AssetStore { // Companies returns the company store associated with this asset store. func (s *AssetStore) Companies() *CompanyStore { return s.company } +// Proxies returns the proxy pool store on the same DB, so agent tool wiring that +// already holds an AssetStore can reach the pool without a separate handle. +func (s *AssetStore) Proxies() *ProxyStore { return s.db.Proxies() } + // withCompanyScopeMutation serializes scope resolution and every asset write // that consumes its result in one transaction. Nested asset side effects reuse // the same transaction through the scoped store. diff --git a/db/proxies.go b/db/proxies.go new file mode 100644 index 0000000..603a78b --- /dev/null +++ b/db/proxies.go @@ -0,0 +1,539 @@ +package db + +import ( + "database/sql" + "encoding/json" + "errors" + "fmt" + "hash/fnv" + "net" + "net/url" + "sort" + "strconv" + "strings" +) + +// ===================================================================== +// 代理池（出口代理轮换） +// ===================================================================== + +// Proxy is one row of the proxies table. Connection info is stored split (not a +// full URL) so entries dedup by (protocol,host,port), the UI can mask the +// password, and logs never carry credentials. Use URL() to build the dial URL. +type Proxy struct { + ID int64 `json:"id"` + Protocol string `json:"protocol"` + Host string `json:"host"` + Port int `json:"port"` + Username string `json:"username,omitempty"` + Password string `json:"password,omitempty"` + Anonymity string `json:"anonymity,omitempty"` + Region string `json:"region,omitempty"` + Tags []string `json:"tags"` + Label string `json:"label,omitempty"` + Source string `json:"source"` + Trusted bool `json:"trusted"` + Enabled bool `json:"enabled"` + Healthy bool `json:"healthy"` + LatencyMs int `json:"latency_ms"` + LastCheckAt string `json:"last_check_at,omitempty"` + LastOkAt string `json:"last_ok_at,omitempty"` + LastError string `json:"last_error,omitempty"` + FailStreak int `json:"fail_streak"` + CheckCount int `json:"check_count"` + OkCount int `json:"ok_count"` + CreatedAt string `json:"created_at,omitempty"` + UpdatedAt string `json:"updated_at,omitempty"` +} + +// URL builds the dial URL from the split fields. Callers use this at connect time +// only — it is never persisted or logged. +func (p *Proxy) URL() *url.URL { + u := &url.URL{Scheme: p.Protocol, Host: net.JoinHostPort(p.Host, strconv.Itoa(p.Port))} + if p.Username != "" || p.Password != "" { + u.User = url.UserPassword(p.Username, p.Password) + } + return u +} + +// Masked returns a copy safe to serialize to the frontend: the password is +// replaced with a fixed placeholder when set, so it is never sent over the wire. +func (p Proxy) Masked() Proxy { + if p.Password != "" { + p.Password = "********" + } + return p +} + +// ProxyFailAutoDisable is the consecutive-failure threshold past which a proxy is +// automatically disabled, so dead free-pool nodes drop out without manual cleanup. +const ProxyFailAutoDisable = 5 + +// SettingProxyPoolEnabled is the settings key for the pool master switch. Defined +// here so both the server (settings UI) and the agent tool layer read one key. +const SettingProxyPoolEnabled = "proxy_pool_enabled" + +var ErrProxyNotFound = errors.New("proxy not found") + +// ProxyStore operates on the proxies + proxy_sources tables. +type ProxyStore struct{ db *DB } + +// Proxies returns a store bound to this DB. +func (d *DB) Proxies() *ProxyStore { return &ProxyStore{db: d} } + +// PoolEnabled reports whether the proxy pool master switch is on (default off). +// Read from settings so the agent tool layer can gate list_proxies without a +// dependency on the server package. +func (s *ProxyStore) PoolEnabled() bool { return s.db.GetBool(SettingProxyPoolEnabled, false) } + +// ProxyFilter narrows a ListProxies / SelectForHost query. Zero value = no filter. +// Limit/Offset apply only to ListProxies (SelectForHost ignores them); Limit<=0 +// means no paging. +type ProxyFilter struct { + Protocol string + Region string + Anonymity string + Tags []string + OnlyHealthy bool + OnlyEnabled bool + TrustedOnly bool + Limit int + Offset int +} + +// listProxyCols is the shared SELECT column list. tags is read as a JSON array +// text (mirrors db/assets.go's array_to_json convention) rather than a driver +// array type, so no lib/pq dependency is needed. +const listProxyCols = `id, protocol, host, port, username, password, + anonymity, region, array_to_json(tags)::text, label, source, trusted, + enabled, healthy, latency_ms, + COALESCE(last_check_at::text,''), COALESCE(last_ok_at::text,''), last_error, + fail_streak, check_count, ok_count, + created_at::text, updated_at::text` + +func scanProxy(rows *sql.Rows) (*Proxy, error) { + var p Proxy + var tagsJSON string + if err := rows.Scan(&p.ID, &p.Protocol, &p.Host, &p.Port, &p.Username, &p.Password, + &p.Anonymity, &p.Region, &tagsJSON, &p.Label, &p.Source, &p.Trusted, + &p.Enabled, &p.Healthy, &p.LatencyMs, + &p.LastCheckAt, &p.LastOkAt, &p.LastError, + &p.FailStreak, &p.CheckCount, &p.OkCount, + &p.CreatedAt, &p.UpdatedAt); err != nil { + return nil, err + } + p.Tags = parseJSONStringArray(tagsJSON) + return &p, nil +} + +// ListProxies returns proxies matching the filter, newest first. When f.Limit>0 +// the result is a page (LIMIT/OFFSET); otherwise all matching rows are returned. +func (s *ProxyStore) ListProxies(f ProxyFilter) ([]*Proxy, error) { + q := `SELECT ` + listProxyCols + ` FROM proxies` + where, args := proxyWhere(f) + if where != "" { + q += " WHERE " + where + } + q += " ORDER BY id DESC" + if f.Limit > 0 { + args = append(args, f.Limit) + q += fmt.Sprintf(" LIMIT $%d", len(args)) + args = append(args, f.Offset) + q += fmt.Sprintf(" OFFSET $%d", len(args)) + } + rows, err := s.db.Query(q, args...) + if err != nil { + return nil, err + } + defer rows.Close() + out := []*Proxy{} + for rows.Next() { + p, err := scanProxy(rows) + if err != nil { + return nil, err + } + out = append(out, p) + } + return out, rows.Err() +} + +// CountProxies returns how many proxies match the filter (ignoring Limit/Offset), +// for server-side pagination totals. +func (s *ProxyStore) CountProxies(f ProxyFilter) (int, error) { + q := `SELECT count(*) FROM proxies` + where, args := proxyWhere(f) + if where != "" { + q += " WHERE " + where + } + var n int + err := s.db.QueryRow(q, args...).Scan(&n) + return n, err +} + +// proxyWhere builds the shared WHERE clause + args ($1-based) from a filter. +func proxyWhere(f ProxyFilter) (string, []any) { + var conds []string + var args []any + add := func(cond string, val any) { + args = append(args, val) + conds = append(conds, fmt.Sprintf(cond, len(args))) + } + if f.Protocol != "" { + add("protocol = $%d", f.Protocol) + } + if f.Region != "" { + add("region = $%d", f.Region) + } + if f.Anonymity != "" { + add("anonymity = $%d", f.Anonymity) + } + if len(f.Tags) > 0 { + add("tags && $%d", marshalStringArray(f.Tags)) + } + if f.OnlyHealthy { + conds = append(conds, "healthy") + } + if f.OnlyEnabled { + conds = append(conds, "enabled") + } + if f.TrustedOnly { + conds = append(conds, "trusted") + } + return strings.Join(conds, " AND "), args +} + +// GetProxy loads one proxy by id. +func (s *ProxyStore) GetProxy(id int64) (*Proxy, error) { + rows, err := s.db.Query(`SELECT `+listProxyCols+` FROM proxies WHERE id = $1`, id) + if err != nil { + return nil, err + } + defer rows.Close() + if !rows.Next() { + if err := rows.Err(); err != nil { + return nil, err + } + return nil, ErrProxyNotFound + } + return scanProxy(rows) +} + +// CreateProxy inserts a manually-entered proxy and returns its id. Duplicate +// (protocol,host,port) returns the existing row's id (idempotent). +func (s *ProxyStore) CreateProxy(p *Proxy) (int64, error) { + if p.Protocol == "" { + p.Protocol = "http" + } + if p.Host == "" || p.Port <= 0 { + return 0, fmt.Errorf("代理需要 host 和 port") + } + if p.Source == "" { + p.Source = "manual" + } + var id int64 + err := s.db.QueryRow(` +INSERT INTO proxies(protocol, host, port, username, password, anonymity, region, tags, label, source, trusted, enabled) +VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12) +ON CONFLICT (protocol, host, port) DO UPDATE SET updated_at = now() +RETURNING id`, + p.Protocol, p.Host, p.Port, p.Username, p.Password, p.Anonymity, p.Region, + marshalStringArray(p.Tags), p.Label, p.Source, p.Trusted, p.Enabled).Scan(&id) + return id, err +} + +// UpdateProxy updates the user-editable fields of one proxy (not health/quality). +func (s *ProxyStore) UpdateProxy(p *Proxy) error { + res, err := s.db.Exec(` +UPDATE proxies SET protocol=$2, host=$3, port=$4, username=$5, password=$6, + anonymity=$7, region=$8, tags=$9, label=$10, trusted=$11, enabled=$12 +WHERE id=$1`, + p.ID, p.Protocol, p.Host, p.Port, p.Username, p.Password, + p.Anonymity, p.Region, marshalStringArray(p.Tags), p.Label, p.Trusted, p.Enabled) + if err != nil { + return err + } + if n, _ := res.RowsAffected(); n == 0 { + return ErrProxyNotFound + } + return nil +} + +// DeleteProxy removes one proxy by id. +func (s *ProxyStore) DeleteProxy(id int64) error { + res, err := s.db.Exec(`DELETE FROM proxies WHERE id = $1`, id) + if err != nil { + return err + } + if n, _ := res.RowsAffected(); n == 0 { + return ErrProxyNotFound + } + return nil +} + +// SetProxyEnabled toggles a proxy's user enable switch. +func (s *ProxyStore) SetProxyEnabled(id int64, enabled bool) error { + res, err := s.db.Exec(`UPDATE proxies SET enabled=$2 WHERE id=$1`, id, enabled) + if err != nil { + return err + } + if n, _ := res.RowsAffected(); n == 0 { + return ErrProxyNotFound + } + return nil +} + +// ImportProxies parses a batch of proxy lines (scheme://[user:pass@]host:port or +// bare host:port, one per line) and inserts them as trusted, source='import'. +// Returns how many new rows were added (duplicates by (protocol,host,port) are +// skipped). Unparseable lines are collected and returned for user feedback. +func (s *ProxyStore) ImportProxies(lines []string) (added int, invalid []string, err error) { + for _, line := range lines { + line = strings.TrimSpace(line) + if line == "" || strings.HasPrefix(line, "#") { + continue + } + p, perr := ParseProxyLine(line) + if perr != nil { + invalid = append(invalid, line) + continue + } + var id int64 + e := s.db.QueryRow(` +INSERT INTO proxies(protocol, host, port, username, password, source, trusted, enabled) +VALUES ($1,$2,$3,$4,$5,'import',true,true) +ON CONFLICT (protocol, host, port) DO NOTHING +RETURNING id`, p.Protocol, p.Host, p.Port, p.Username, p.Password).Scan(&id) + switch { + case e == sql.ErrNoRows: // duplicate, skipped + case e != nil: + return added, invalid, e + default: + added++ + } + } + return added, invalid, nil +} + +// UpsertFromSource inserts proxies fetched from a free source as untrusted, +// source=. Existing (protocol,host,port) rows are left untouched +// (a manually-added trusted entry is never downgraded). Returns rows added. +func (s *ProxyStore) UpsertFromSource(sourceName string, proxies []Proxy) (int, error) { + added := 0 + for i := range proxies { + p := proxies[i] + if p.Protocol == "" || p.Host == "" || p.Port <= 0 { + continue + } + var id int64 + e := s.db.QueryRow(` +INSERT INTO proxies(protocol, host, port, anonymity, region, source, trusted, enabled) +VALUES ($1,$2,$3,$4,$5,$6,false,true) +ON CONFLICT (protocol, host, port) DO NOTHING +RETURNING id`, p.Protocol, p.Host, p.Port, p.Anonymity, p.Region, sourceName).Scan(&id) + switch { + case e == sql.ErrNoRows: + case e != nil: + return added, e + default: + added++ + } + } + return added, nil +} + +// SelectForHost picks a proxy for a target host with per-host stickiness: the same +// host always maps to the same egress (until the healthy set changes), so a scan's +// requests to one target share an exit IP. Candidates are enabled+healthy (and +// trusted when trustedOnly), ranked by quality (low fail_streak, high success +// rate, recent success); a stable index derived from the host name picks within +// the ranked set. Returns nil when no candidate exists (caller falls back direct). +func (s *ProxyStore) SelectForHost(host string, trustedOnly bool) (*Proxy, error) { + pool, err := s.ListProxies(ProxyFilter{OnlyEnabled: true, OnlyHealthy: true, TrustedOnly: trustedOnly}) + if err != nil { + return nil, err + } + if len(pool) == 0 { + return nil, nil + } + // Quality order: fewer consecutive fails, higher success rate, more recent success. + sort.SliceStable(pool, func(i, j int) bool { + a, b := pool[i], pool[j] + if a.FailStreak != b.FailStreak { + return a.FailStreak < b.FailStreak + } + ra, rb := successRate(a), successRate(b) + if ra != rb { + return ra > rb + } + return a.LastOkAt > b.LastOkAt + }) + // Stable per-host index into the ranked set → same host, same exit. + h := fnv.New32a() + _, _ = h.Write([]byte(strings.ToLower(host))) + return pool[int(h.Sum32())%len(pool)], nil +} + +func successRate(p *Proxy) float64 { + if p.CheckCount == 0 { + return 0 + } + return float64(p.OkCount) / float64(p.CheckCount) +} + +// UpdateHealth records the outcome of one probe. +// - success: clears the fail streak, stamps last_ok_at, marks healthy. +// - failure of an UNTRUSTED (free-source) proxy: the row is DELETED outright — +// free proxies are disposable and re-fetched, so a dead one is just removed. +// - failure of a TRUSTED (manual/imported) proxy: kept for retry; the fail streak +// increments and past ProxyFailAutoDisable the proxy auto-disables (not deleted, +// since the user entered it deliberately). +func (s *ProxyStore) UpdateHealth(id int64, ok bool, latencyMs int, probeErr string) error { + if ok { + _, err := s.db.Exec(` +UPDATE proxies SET healthy=true, latency_ms=$2, last_error='', + last_check_at=now(), last_ok_at=now(), fail_streak=0, + check_count=check_count+1, ok_count=ok_count+1 +WHERE id=$1`, id, latencyMs) + return err + } + // Free-source (untrusted) proxy failed → delete it. RowsAffected>0 means it was + // an untrusted row and is now gone; nothing more to do. + res, err := s.db.Exec(`DELETE FROM proxies WHERE id=$1 AND NOT trusted`, id) + if err != nil { + return err + } + if n, _ := res.RowsAffected(); n > 0 { + return nil + } + // Trusted proxy failed → keep, bump streak, auto-disable past the threshold. + _, err = s.db.Exec(` +UPDATE proxies SET healthy=false, last_error=$2, last_check_at=now(), + fail_streak=fail_streak+1, check_count=check_count+1, + enabled = CASE WHEN fail_streak+1 >= $3 THEN false ELSE enabled END +WHERE id=$1`, id, probeErr, ProxyFailAutoDisable) + return err +} + +// ProxySource is one free-pool source's enable switch + last fetch status. +type ProxySource struct { + Name string `json:"name"` + Enabled bool `json:"enabled"` + LastFetchAt string `json:"last_fetch_at,omitempty"` + LastCount int `json:"last_count"` + LastError string `json:"last_error,omitempty"` +} + +// ListSources returns the persisted state of every source name passed in, filling +// defaults (disabled, never fetched) for names with no row yet. This keeps the +// code-defined source catalog as the source of truth for which sources exist. +func (s *ProxyStore) ListSources(names []string) ([]ProxySource, error) { + rows, err := s.db.Query(`SELECT name, enabled, COALESCE(last_fetch_at::text,''), last_count, last_error FROM proxy_sources`) + if err != nil { + return nil, err + } + defer rows.Close() + byName := map[string]ProxySource{} + for rows.Next() { + var src ProxySource + if err := rows.Scan(&src.Name, &src.Enabled, &src.LastFetchAt, &src.LastCount, &src.LastError); err != nil { + return nil, err + } + byName[src.Name] = src + } + if err := rows.Err(); err != nil { + return nil, err + } + out := make([]ProxySource, 0, len(names)) + for _, name := range names { + if src, ok := byName[name]; ok { + out = append(out, src) + } else { + out = append(out, ProxySource{Name: name}) + } + } + return out, nil +} + +// SetSourceEnabled toggles one free-pool source (upserts the row). +func (s *ProxyStore) SetSourceEnabled(name string, enabled bool) error { + _, err := s.db.Exec(` +INSERT INTO proxy_sources(name, enabled) VALUES ($1,$2) +ON CONFLICT (name) DO UPDATE SET enabled = EXCLUDED.enabled`, name, enabled) + return err +} + +// EnabledSources returns the names of sources currently switched on. +func (s *ProxyStore) EnabledSources() ([]string, error) { + rows, err := s.db.Query(`SELECT name FROM proxy_sources WHERE enabled`) + if err != nil { + return nil, err + } + defer rows.Close() + var out []string + for rows.Next() { + var name string + if err := rows.Scan(&name); err != nil { + return nil, err + } + out = append(out, name) + } + return out, rows.Err() +} + +// RecordFetch stamps a source's last fetch outcome. +func (s *ProxyStore) RecordFetch(name string, count int, fetchErr string) error { + _, err := s.db.Exec(` +INSERT INTO proxy_sources(name, enabled, last_fetch_at, last_count, last_error) +VALUES ($1, true, now(), $2, $3) +ON CONFLICT (name) DO UPDATE SET last_fetch_at = now(), last_count = $2, last_error = $3`, + name, count, fetchErr) + return err +} + +// ParseProxyLine parses "scheme://[user:pass@]host:port" or a bare "host:port" +// (defaulting to http) into a Proxy's connection fields. +func ParseProxyLine(line string) (Proxy, error) { + line = strings.TrimSpace(line) + if line == "" { + return Proxy{}, fmt.Errorf("空行") + } + if !strings.Contains(line, "://") { + line = "http://" + line + } + u, err := url.Parse(line) + if err != nil { + return Proxy{}, err + } + scheme := strings.ToLower(u.Scheme) + switch scheme { + case "http", "https", "socks5": + default: + return Proxy{}, fmt.Errorf("不支持的协议: %s（仅支持 http/https/socks5）", u.Scheme) + } + host := u.Hostname() + portStr := u.Port() + if host == "" || portStr == "" { + return Proxy{}, fmt.Errorf("需要 host:port") + } + port, err := strconv.Atoi(portStr) + if err != nil || port <= 0 || port > 65535 { + return Proxy{}, fmt.Errorf("无效端口: %s", portStr) + } + p := Proxy{Protocol: scheme, Host: host, Port: port} + if u.User != nil { + p.Username = u.User.Username() + p.Password, _ = u.User.Password() + } + return p, nil +} + +// parseJSONStringArray unmarshals an array_to_json text into a []string, matching +// how db/assets.go reads TEXT[] columns. Always returns non-nil for a valid array. +func parseJSONStringArray(jsonText string) []string { + out := []string{} + if jsonText == "" || jsonText == "null" { + return out + } + _ = json.Unmarshal([]byte(jsonText), &out) + return out +} diff --git a/db/proxies_test.go b/db/proxies_test.go new file mode 100644 index 0000000..61d7406 --- /dev/null +++ b/db/proxies_test.go @@ -0,0 +1,229 @@ +package db + +import ( + "fmt" + "testing" + "time" +) + +// TestParseProxyLine covers URL/host:port parsing + protocol validation, no DB. +func TestParseProxyLine(t *testing.T) { + cases := []struct { + in string + proto string + host string + port int + user string + pass string + wantErr bool + }{ + {in: "1.2.3.4:8080", proto: "http", host: "1.2.3.4", port: 8080}, + {in: "socks5://9.9.9.9:1080", proto: "socks5", host: "9.9.9.9", port: 1080}, + {in: "http://user:pass@10.0.0.1:3128", proto: "http", host: "10.0.0.1", port: 3128, user: "user", pass: "pass"}, + {in: "socks4://5.6.7.8:1080", wantErr: true}, // socks4 不支持（仅 http/https/socks5） + {in: " 8.8.8.8:80 ", proto: "http", host: "8.8.8.8", port: 80}, + {in: "ftp://1.2.3.4:21", wantErr: true}, // unsupported scheme + {in: "1.2.3.4", wantErr: true}, // no port + {in: "1.2.3.4:notaport", wantErr: true}, // bad port + {in: "1.2.3.4:99999", wantErr: true}, // out of range + {in: "", wantErr: true}, + } + for _, c := range cases { + p, err := ParseProxyLine(c.in) + if c.wantErr { + if err == nil { + t.Errorf("ParseProxyLine(%q) want error, got %+v", c.in, p) + } + continue + } + if err != nil { + t.Errorf("ParseProxyLine(%q) unexpected error: %v", c.in, err) + continue + } + if p.Protocol != c.proto || p.Host != c.host || p.Port != c.port || p.Username != c.user || p.Password != c.pass { + t.Errorf("ParseProxyLine(%q) = %+v, want proto=%s host=%s port=%d user=%s pass=%s", + c.in, p, c.proto, c.host, c.port, c.user, c.pass) + } + } +} + +// TestProxyURLRoundTrip verifies URL() rebuilds a dial URL including auth. +func TestProxyURLRoundTrip(t *testing.T) { + p := Proxy{Protocol: "socks5", Host: "1.2.3.4", Port: 1080, Username: "u", Password: "p"} + if got := p.URL().String(); got != "socks5://u:p@1.2.3.4:1080" { + t.Fatalf("URL = %q", got) + } + p2 := Proxy{Protocol: "http", Host: "5.6.7.8", Port: 8080} + if got := p2.URL().String(); got != "http://5.6.7.8:8080" { + t.Fatalf("URL = %q", got) + } +} + +// TestProxyMasked ensures the password is never serialized in the clear. +func TestProxyMasked(t *testing.T) { + p := Proxy{Password: "secret"} + if p.Masked().Password != "********" { + t.Fatalf("password not masked: %q", p.Masked().Password) + } + if (Proxy{}).Masked().Password != "" { + t.Fatalf("empty password should stay empty") + } +} + +// TestProxyStoreLifecycle exercises import de-dup, health update + auto-disable, +// per-host sticky selection, and trusted-only filtering against dev PG. +func TestProxyStoreLifecycle(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + ps := d.Proxies() + + // Unique host octet per run so parallel/repeat runs don't collide. + seed := time.Now().UnixNano() % 250 + host := func(n int64) string { return fmt.Sprintf("203.0.113.%d", (seed+n)%254+1) } + var ids []int64 + t.Cleanup(func() { + for _, id := range ids { + _ = ps.DeleteProxy(id) + } + }) + + // Import three trusted proxies; the duplicate line must be skipped. + lines := []string{ + "http://" + host(0) + ":8080", + "http://" + host(0) + ":8080", // dup + "socks5://" + host(1) + ":1080", + } + added, invalid, err := ps.ImportProxies(lines) + if err != nil { + t.Fatal(err) + } + if added != 2 || len(invalid) != 0 { + t.Fatalf("import added=%d invalid=%v, want added=2", added, invalid) + } + + all, err := ps.ListProxies(ProxyFilter{}) + if err != nil { + t.Fatal(err) + } + for _, p := range all { + if p.Host == host(0) || p.Host == host(1) { + ids = append(ids, p.ID) + } + } + if len(ids) != 2 { + t.Fatalf("want 2 imported rows, got %d", len(ids)) + } + + // Before any probe, nothing is healthy → no selection. + pick, err := ps.SelectForHost("target.example.com", true) + if err != nil { + t.Fatal(err) + } + if pick != nil { + t.Fatalf("no healthy proxy yet, got %+v", pick) + } + + // Mark both healthy, then selection must return one and be sticky per host. + for _, id := range ids { + if err := ps.UpdateHealth(id, true, 100, ""); err != nil { + t.Fatal(err) + } + } + p1, err := ps.SelectForHost("stickyhost", true) + if err != nil || p1 == nil { + t.Fatalf("select: %v %+v", err, p1) + } + p2, _ := ps.SelectForHost("stickyhost", true) + if p2 == nil || p1.ID != p2.ID { + t.Fatalf("per-host stickiness broken: %v vs %v", p1, p2) + } + + // trustedOnly=false must still return (all imports are trusted anyway). + if p, _ := ps.SelectForHost("x", false); p == nil { + t.Fatal("select trustedOnly=false returned nil") + } + + // Fail one proxy ProxyFailAutoDisable times → auto-disabled → drops out. + victim := ids[0] + for i := 0; i < ProxyFailAutoDisable; i++ { + if err := ps.UpdateHealth(victim, false, 0, "timeout"); err != nil { + t.Fatal(err) + } + } + got, err := ps.GetProxy(victim) + if err != nil { + t.Fatal(err) + } + if got.Enabled { + t.Fatalf("proxy should auto-disable after %d fails", ProxyFailAutoDisable) + } + + // Free-source (untrusted) proxy: a single probe failure deletes it outright. + freeAdded, err := ps.UpsertFromSource("unittest-free", []Proxy{{Protocol: "http", Host: host(2), Port: 8080}}) + if err != nil || freeAdded != 1 { + t.Fatalf("seed free proxy: added=%d err=%v", freeAdded, err) + } + all, err = ps.ListProxies(ProxyFilter{}) + if err != nil { + t.Fatal(err) + } + var freeID int64 + for _, p := range all { + if p.Host == host(2) { + freeID = p.ID + ids = append(ids, p.ID) + } + } + if freeID == 0 { + t.Fatal("free proxy not found after upsert") + } + if err := ps.UpdateHealth(freeID, false, 0, "timeout"); err != nil { + t.Fatal(err) + } + if _, err := ps.GetProxy(freeID); err != ErrProxyNotFound { + t.Fatalf("free-source proxy should be deleted on probe failure, got err=%v", err) + } +} + +// TestProxySourceToggle covers source enable persistence + fetch stamping. +func TestProxySourceToggle(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + ps := d.Proxies() + + name := fmt.Sprintf("unittest-source-%d", time.Now().UnixNano()) + t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM proxy_sources WHERE name=$1`, name) }) + + if err := ps.SetSourceEnabled(name, true); err != nil { + t.Fatal(err) + } + enabled, err := ps.EnabledSources() + if err != nil { + t.Fatal(err) + } + found := false + for _, n := range enabled { + if n == name { + found = true + } + } + if !found { + t.Fatalf("source %s not in enabled set %v", name, enabled) + } + if err := ps.RecordFetch(name, 42, ""); err != nil { + t.Fatal(err) + } + sources, err := ps.ListSources([]string{name}) + if err != nil || len(sources) != 1 { + t.Fatalf("list sources: %v %+v", err, sources) + } + if sources[0].LastCount != 42 || !sources[0].Enabled { + t.Fatalf("fetch not recorded: %+v", sources[0]) + } +} diff --git a/db/schema.sql b/db/schema.sql index 968bd74..bf92180 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -782,3 +782,50 @@ CREATE TABLE IF NOT EXISTS server_logs ( text TEXT NOT NULL DEFAULT '' ); CREATE INDEX IF NOT EXISTS idx_server_logs_id ON server_logs(id DESC); + +-- ===================================================================== +-- N. 代理池（出口代理轮换） +-- ===================================================================== +-- 连接信息拆分存（非完整 URL）：去重键 (protocol,host,port)、UI 隐藏密码、日志不落认证。 +-- username/password 明文存（与 llm_profiles.api_key 一致）。source/trusted 供安全阀门， +-- 免费源抓来的 trusted=false，默认不进主出口轮换。质量字段驱动自动禁用死代理 + 优先稳定节点。 +CREATE TABLE IF NOT EXISTS proxies ( + id BIGSERIAL PRIMARY KEY, + protocol TEXT NOT NULL DEFAULT 'http', -- http/https/socks5 + host TEXT NOT NULL, + port INTEGER NOT NULL, + username TEXT NOT NULL DEFAULT '', + password TEXT NOT NULL DEFAULT '', + anonymity TEXT NOT NULL DEFAULT '', -- elite/anonymous/transparent/'' + region TEXT NOT NULL DEFAULT '', -- 国家码 CN/US… + tags TEXT[] NOT NULL DEFAULT '{}', + label TEXT NOT NULL DEFAULT '', + source TEXT NOT NULL DEFAULT 'manual', -- manual/import/<源名> + trusted BOOLEAN NOT NULL DEFAULT true, + enabled BOOLEAN NOT NULL DEFAULT true, + healthy BOOLEAN NOT NULL DEFAULT false, + latency_ms INTEGER NOT NULL DEFAULT 0, + last_check_at TIMESTAMPTZ, + last_ok_at TIMESTAMPTZ, + last_error TEXT NOT NULL DEFAULT '', + fail_streak INTEGER NOT NULL DEFAULT 0, + check_count INTEGER NOT NULL DEFAULT 0, + ok_count INTEGER NOT NULL DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE(protocol, host, port) +); +CREATE INDEX IF NOT EXISTS idx_proxies_tags ON proxies USING GIN(tags); +CREATE INDEX IF NOT EXISTS idx_proxies_pick ON proxies(enabled, healthy, trusted); +DROP TRIGGER IF EXISTS trg_proxies_upd ON proxies; +CREATE TRIGGER trg_proxies_upd BEFORE UPDATE ON proxies + FOR EACH ROW EXECUTE FUNCTION set_updated_at(); + +-- 免费代理源开关与抓取状态。默认无行 = 全部未开启（安全底线）。 +CREATE TABLE IF NOT EXISTS proxy_sources ( + name TEXT PRIMARY KEY, + enabled BOOLEAN NOT NULL DEFAULT false, + last_fetch_at TIMESTAMPTZ, + last_count INTEGER NOT NULL DEFAULT 0, + last_error TEXT NOT NULL DEFAULT '' +); diff --git a/proxypool/connect.go b/proxypool/connect.go new file mode 100644 index 0000000..f97f97d --- /dev/null +++ b/proxypool/connect.go @@ -0,0 +1,69 @@ +package proxypool + +import ( + "bufio" + "context" + "crypto/tls" + "encoding/base64" + "fmt" + "net" + "net/http" + "net/url" + "time" +) + +// dialHTTPConnect tunnels to address through an http/https proxy via a CONNECT +// request. For an https proxy the hop to the proxy itself is wrapped in TLS +// first. Mirrors the well-worn net/http dialConn CONNECT flow. +func dialHTTPConnect(ctx context.Context, proxyURL *url.URL, address string) (net.Conn, error) { + d := &net.Dialer{} + conn, err := d.DialContext(ctx, "tcp", proxyURL.Host) + if err != nil { + return nil, err + } + if proxyURL.Scheme == "https" { + tlsConn := tls.Client(conn, &tls.Config{ServerName: proxyURL.Hostname()}) + if err := tlsConn.HandshakeContext(ctx); err != nil { + conn.Close() + return nil, err + } + conn = tlsConn + } + req := &http.Request{ + Method: http.MethodConnect, + URL: &url.URL{Opaque: address}, + Host: address, + Header: http.Header{}, + } + if proxyURL.User != nil { + req.Header.Set("Proxy-Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(proxyURL.User.String()))) + } + connectCtx, cancel := context.WithTimeout(ctx, time.Minute) + defer cancel() + done := make(chan error, 1) + var resp *http.Response + go func() { + if werr := req.Write(conn); werr != nil { + done <- werr + return + } + r, rerr := http.ReadResponse(bufio.NewReader(conn), req) + resp = r + done <- rerr + }() + select { + case <-connectCtx.Done(): + conn.Close() + return nil, connectCtx.Err() + case err = <-done: + } + if err != nil { + conn.Close() + return nil, err + } + if resp.StatusCode != http.StatusOK { + conn.Close() + return nil, fmt.Errorf("proxy CONNECT failed: %s", resp.Status) + } + return conn, nil +} diff --git a/proxypool/dial.go b/proxypool/dial.go new file mode 100644 index 0000000..b37f23a --- /dev/null +++ b/proxypool/dial.go @@ -0,0 +1,47 @@ +// Package proxypool provides the outbound proxy pool: dialing through a proxy of +// any supported scheme, active health probing, and fetching free proxy sources. +// It is deliberately independent of the traffic (MITM) layer so the pool works +// whether or not traffic recording is on. +package proxypool + +import ( + "context" + "fmt" + "net" + "net/url" + "strings" + + xproxy "golang.org/x/net/proxy" +) + +// DialThrough opens a TCP connection to address (host:port) through the given +// proxy URL. Supports http/https (CONNECT tunnel) and socks5. The returned +// conn is a raw byte stream to the target; TLS to the target (if any) is the +// caller's job — the proxy only tunnels, it never terminates TLS to the target. +func DialThrough(ctx context.Context, proxyURL *url.URL, address string) (net.Conn, error) { + switch strings.ToLower(proxyURL.Scheme) { + case "socks5": + return dialSOCKS(ctx, proxyURL, address) + case "http", "https": + return dialHTTPConnect(ctx, proxyURL, address) + default: + return nil, fmt.Errorf("unsupported proxy scheme %q", proxyURL.Scheme) + } +} + +// dialSOCKS tunnels through a SOCKS5 proxy. +func dialSOCKS(ctx context.Context, proxyURL *url.URL, address string) (net.Conn, error) { + var auth *xproxy.Auth + if proxyURL.User != nil { + pass, _ := proxyURL.User.Password() + auth = &xproxy.Auth{User: proxyURL.User.Username(), Password: pass} + } + dialer, err := xproxy.SOCKS5("tcp", proxyURL.Host, auth, xproxy.Direct) + if err != nil { + return nil, err + } + if cd, ok := dialer.(xproxy.ContextDialer); ok { + return cd.DialContext(ctx, "tcp", address) + } + return dialer.Dial("tcp", address) +} diff --git a/proxypool/gateway.go b/proxypool/gateway.go new file mode 100644 index 0000000..15c6f90 --- /dev/null +++ b/proxypool/gateway.go @@ -0,0 +1,174 @@ +package proxypool + +import ( + "context" + "errors" + "io" + "log" + "net" + "net/http" + "net/url" + "sync/atomic" + "time" +) + +// Gateway is a local forward proxy (入口C): agents point HTTP_PROXY at it and it +// tunnels each connection out through a pool proxy chosen per target host. It is a +// plain CONNECT tunnel / HTTP forwarder — it does NOT decrypt TLS or record, so it +// works with cert-pinned targets and protocols the MITM can't parse. Independent +// of the traffic layer; used when traffic capture is off but the pool is on. +type Gateway struct { + addr string + srv *http.Server + upstream atomic.Pointer[func(host string) *url.URL] +} + +// NewGateway builds a gateway listening on addr (e.g. 127.0.0.1:8789). +func NewGateway(addr string) *Gateway { return &Gateway{addr: addr} } + +// Addr returns the listen address. +func (g *Gateway) Addr() string { return g.addr } + +// SetUpstream installs (or clears with nil) the per-host upstream resolver. A nil +// return for a host means dial that host directly (gateway adds nothing). +func (g *Gateway) SetUpstream(fn func(host string) *url.URL) { + if fn == nil { + g.upstream.Store(nil) + return + } + g.upstream.Store(&fn) +} + +func (g *Gateway) pick(host string) *url.URL { + if fn := g.upstream.Load(); fn != nil { + return (*fn)(hostOnlyGW(host)) + } + return nil +} + +// Start begins serving in the background. Safe to call once. +func (g *Gateway) Start() error { + ln, err := net.Listen("tcp", g.addr) + if err != nil { + return err + } + g.srv = &http.Server{Handler: http.HandlerFunc(g.handle)} + go func() { + if err := g.srv.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) { + log.Printf("[proxygw] stopped: %v", err) + } + }() + log.Printf("[proxypool] gateway on %s", g.addr) + return nil +} + +// Stop gracefully shuts the gateway down. +func (g *Gateway) Stop() { + if g.srv != nil { + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + _ = g.srv.Shutdown(ctx) + } +} + +func (g *Gateway) handle(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodConnect { + g.handleConnect(w, r) + return + } + g.handleHTTP(w, r) +} + +// handleConnect tunnels an HTTPS (or any TCP) target: dial out (through the chosen +// pool proxy, or direct), 200 the client, then splice the two connections. +func (g *Gateway) handleConnect(w http.ResponseWriter, r *http.Request) { + target := r.Host // host:port + ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second) + defer cancel() + upConn, err := g.dialOut(ctx, target) + if err != nil { + http.Error(w, "gateway dial failed: "+err.Error(), http.StatusBadGateway) + return + } + hj, ok := w.(http.Hijacker) + if !ok { + upConn.Close() + http.Error(w, "hijack unsupported", http.StatusInternalServerError) + return + } + clientConn, _, err := hj.Hijack() + if err != nil { + upConn.Close() + return + } + if _, err := clientConn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n")); err != nil { + clientConn.Close() + upConn.Close() + return + } + splice(clientConn, upConn) +} + +// handleHTTP forwards a plain (non-CONNECT) HTTP request through the chosen pool +// proxy (or direct) and copies the response back. +func (g *Gateway) handleHTTP(w http.ResponseWriter, r *http.Request) { + tr := &http.Transport{DisableKeepAlives: true} + if up := g.pick(r.Host); up != nil { + switch up.Scheme { + case "http", "https": + tr.Proxy = http.ProxyURL(up) + case "socks5": + tr.DialContext = func(ctx context.Context, _, address string) (net.Conn, error) { + return dialSOCKS(ctx, up, address) + } + } + } + defer tr.CloseIdleConnections() + outReq := r.Clone(r.Context()) + outReq.RequestURI = "" + resp, err := tr.RoundTrip(outReq) + if err != nil { + http.Error(w, "gateway forward failed: "+err.Error(), http.StatusBadGateway) + return + } + defer resp.Body.Close() + for k, vs := range resp.Header { + for _, v := range vs { + w.Header().Add(k, v) + } + } + w.WriteHeader(resp.StatusCode) + _, _ = io.Copy(w, resp.Body) +} + +// dialOut opens a raw connection to target, through the pool proxy chosen for its +// host when one is available, else directly. +func (g *Gateway) dialOut(ctx context.Context, target string) (net.Conn, error) { + if up := g.pick(target); up != nil { + return DialThrough(ctx, up, target) + } + var d net.Dialer + return d.DialContext(ctx, "tcp", target) +} + +// splice copies bytes both ways between two connections until either side closes. +func splice(a, b net.Conn) { + done := make(chan struct{}, 2) + cp := func(dst, src net.Conn) { + _, _ = io.Copy(dst, src) + done <- struct{}{} + } + go cp(a, b) + go cp(b, a) + <-done + a.Close() + b.Close() +} + +// hostOnlyGW strips a trailing :port so the upstream resolver keys on bare host. +func hostOnlyGW(hostport string) string { + if h, _, err := net.SplitHostPort(hostport); err == nil { + return h + } + return hostport +} diff --git a/proxypool/gateway_test.go b/proxypool/gateway_test.go new file mode 100644 index 0000000..a7826a4 --- /dev/null +++ b/proxypool/gateway_test.go @@ -0,0 +1,87 @@ +package proxypool + +import ( + "bufio" + "io" + "net" + "net/http" + "net/url" + "sync/atomic" + "testing" + "time" +) + +// TestGatewayConnectTunnel starts the gateway in direct mode (no upstream) and +// verifies a CONNECT tunnel to a local echo server round-trips bytes. +func TestGatewayConnectTunnel(t *testing.T) { + // Echo target: writes back whatever it reads. + target, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer target.Close() + go func() { + for { + c, err := target.Accept() + if err != nil { + return + } + go func(c net.Conn) { _, _ = io.Copy(c, c); c.Close() }(c) + } + }() + + gw := NewGateway("127.0.0.1:0") + // Bind an explicit listener so we know the port (NewGateway+Start uses g.addr). + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + var picked atomic.Int32 + gw.SetUpstream(func(string) *url.URL { picked.Add(1); return nil }) // direct + gw.srv = &http.Server{Handler: http.HandlerFunc(gw.handle)} + go func() { _ = gw.srv.Serve(ln) }() + defer gw.Stop() + + // Client dials the gateway, sends CONNECT to the echo target, then echoes. + conn, err := net.DialTimeout("tcp", ln.Addr().String(), 2*time.Second) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + if _, err := conn.Write([]byte("CONNECT " + target.Addr().String() + " HTTP/1.1\r\nHost: " + target.Addr().String() + "\r\n\r\n")); err != nil { + t.Fatal(err) + } + br := bufio.NewReader(conn) + status, err := br.ReadString('\n') + if err != nil || status != "HTTP/1.1 200 Connection Established\r\n" { + t.Fatalf("connect status=%q err=%v", status, err) + } + // Consume the blank line terminating the CONNECT response headers. + if _, err := br.ReadString('\n'); err != nil { + t.Fatal(err) + } + if _, err := conn.Write([]byte("ping")); err != nil { + t.Fatal(err) + } + buf := make([]byte, 4) + if _, err := io.ReadFull(br, buf); err != nil { + t.Fatal(err) + } + if string(buf) != "ping" { + t.Fatalf("echo = %q, want ping", string(buf)) + } + if picked.Load() == 0 { + t.Fatal("upstream resolver was not consulted") + } +} + +// TestGatewayPickUsesHostOnly verifies the resolver receives a bare host (no port). +func TestGatewayPickUsesHostOnly(t *testing.T) { + gw := NewGateway("127.0.0.1:0") + var gotHost string + gw.SetUpstream(func(host string) *url.URL { gotHost = host; return nil }) + _ = gw.pick("example.com:443") + if gotHost != "example.com" { + t.Fatalf("resolver host = %q, want example.com", gotHost) + } +} diff --git a/proxypool/pool.go b/proxypool/pool.go new file mode 100644 index 0000000..08ecd64 --- /dev/null +++ b/proxypool/pool.go @@ -0,0 +1,205 @@ +package proxypool + +import ( + "context" + "errors" + "log" + "net/http" + "sync" + "time" + + "github.com/Autumn-27/artex/db" +) + +// errUnknownSource is returned when a fetch targets a name not in BuiltinSources. +var errUnknownSource = errors.New("unknown proxy source") + +// Config wires the pool's background loops to live settings via callbacks, so a +// settings change takes effect on the next tick without restarting the loops. +type Config struct { + Enabled func() bool // master switch (proxy_pool_enabled) + FetchInterval func() time.Duration // how often to pull enabled free sources + CheckInterval func() time.Duration // how often to re-probe every enabled proxy + ProbeURL func() string // liveness target (empty = DefaultProbeURL) + Concurrency int // parallel probes (default 50) + ProbeTimeout time.Duration // per-probe timeout (default 10s) +} + +// Pool runs the proxy pool's two background loops (fetch + probe) and serves +// one-off manual probes. It owns nothing the DB doesn't; all state lives in PG. +type Pool struct { + db *db.DB + cfg Config + client *http.Client // for fetching source lists (direct, not through the pool) + cancel context.CancelFunc + wg sync.WaitGroup +} + +// NewPool builds a pool bound to db with the given config, applying defaults. +func NewPool(database *db.DB, cfg Config) *Pool { + if cfg.Concurrency <= 0 { + cfg.Concurrency = 50 + } + if cfg.ProbeTimeout <= 0 { + cfg.ProbeTimeout = 10 * time.Second + } + return &Pool{ + db: database, + cfg: cfg, + client: &http.Client{Timeout: 30 * time.Second}, + } +} + +// Start launches the fetch and probe loops. Idempotent guards are the caller's +// job; call once. Stop() (or a cancelled parent ctx) ends both loops. +func (p *Pool) Start(parent context.Context) { + ctx, cancel := context.WithCancel(parent) + p.cancel = cancel + p.wg.Add(2) + go p.loop(ctx, p.cfg.FetchInterval, 15*time.Minute, p.fetchOnce) + go p.loop(ctx, p.cfg.CheckInterval, 30*time.Minute, p.probeOnce) +} + +// Stop ends the background loops and waits for them to exit. +func (p *Pool) Stop() { + if p.cancel != nil { + p.cancel() + } + p.wg.Wait() +} + +// loop runs work on a self-rescheduling timer whose interval is re-read each round +// (so settings changes apply next tick). It skips work while the pool is disabled. +func (p *Pool) loop(ctx context.Context, interval func() time.Duration, fallback time.Duration, work func(context.Context)) { + defer p.wg.Done() + next := func() time.Duration { + d := fallback + if interval != nil { + if v := interval(); v > 0 { + d = v + } + } + if d < time.Minute { + d = time.Minute // floor: never hammer sources/targets + } + return d + } + timer := time.NewTimer(next()) + defer timer.Stop() + for { + select { + case <-ctx.Done(): + return + case <-timer.C: + if p.enabled() { + work(ctx) + } + timer.Reset(next()) + } + } +} + +func (p *Pool) enabled() bool { return p.cfg.Enabled == nil || p.cfg.Enabled() } + +func (p *Pool) probeURL() string { + if p.cfg.ProbeURL == nil { + return DefaultProbeURL + } + return p.cfg.ProbeURL() +} + +// fetchOnce pulls every enabled free source and upserts new (untrusted) proxies, +// then immediately probes the freshly-added ones so they don't sit unknown for a +// full check interval. +func (p *Pool) fetchOnce(ctx context.Context) { + names, err := p.db.Proxies().EnabledSources() + if err != nil { + log.Printf("[proxypool] list enabled sources: %v", err) + return + } + for _, name := range names { + _, _, _ = p.fetchSource(ctx, name) + } + // Probe whatever is now enabled+unknown so new nodes become usable fast. + p.probeOnce(ctx) +} + +// fetchSource pulls one source by name and upserts its proxies. Returns the total +// fetched and how many were newly added. Records the outcome regardless of the +// source's enabled state, so the manual "fetch now" button works on any source. +func (p *Pool) fetchSource(ctx context.Context, name string) (total, added int, err error) { + store := p.db.Proxies() + src, ok := sourceByName(name) + if !ok { + return 0, 0, errUnknownSource + } + proxies, ferr := fetch(ctx, src, p.client) + if ferr != nil { + _ = store.RecordFetch(name, 0, trimErr(ferr.Error())) + return 0, 0, ferr + } + added, uerr := store.UpsertFromSource(name, proxies) + if uerr != nil { + _ = store.RecordFetch(name, len(proxies), trimErr(uerr.Error())) + return len(proxies), 0, uerr + } + _ = store.RecordFetch(name, len(proxies), "") + log.Printf("[proxypool] source %s: fetched %d, added %d new", name, len(proxies), added) + return len(proxies), added, nil +} + +// FetchSourceNow pulls one source on demand (manual "fetch now" button), then +// probes so freshly-added nodes become usable without waiting for the loop. +func (p *Pool) FetchSourceNow(ctx context.Context, name string) (total, added int, err error) { + total, added, err = p.fetchSource(ctx, name) + if err == nil && added > 0 { + go p.probeOnce(context.WithoutCancel(ctx)) + } + return total, added, err +} + +// probeOnce concurrently probes every enabled proxy and writes back health. +func (p *Pool) probeOnce(ctx context.Context) { + store := p.db.Proxies() + proxies, err := store.ListProxies(db.ProxyFilter{OnlyEnabled: true}) + if err != nil { + log.Printf("[proxypool] list for probe: %v", err) + return + } + if len(proxies) == 0 { + return + } + sem := make(chan struct{}, p.cfg.Concurrency) + var wg sync.WaitGroup + for _, pr := range proxies { + select { + case <-ctx.Done(): + wg.Wait() + return + case sem <- struct{}{}: + } + wg.Add(1) + go func(pr *db.Proxy) { + defer wg.Done() + defer func() { <-sem }() + res := Probe(ctx, pr.URL(), p.probeURL(), p.cfg.ProbeTimeout) + _ = store.UpdateHealth(pr.ID, res.OK, int(res.Latency.Milliseconds()), res.Err) + }(pr) + } + wg.Wait() +} + +// ProbeNow probes a single proxy on demand (manual "check" button) and writes back +// the result, returning it for immediate UI feedback. +func (p *Pool) ProbeNow(ctx context.Context, id int64) (ProbeResult, error) { + store := p.db.Proxies() + pr, err := store.GetProxy(id) + if err != nil { + return ProbeResult{}, err + } + res := Probe(ctx, pr.URL(), p.probeURL(), p.cfg.ProbeTimeout) + if err := store.UpdateHealth(id, res.OK, int(res.Latency.Milliseconds()), res.Err); err != nil { + return res, err + } + return res, nil +} diff --git a/proxypool/probe.go b/proxypool/probe.go new file mode 100644 index 0000000..e8c04f3 --- /dev/null +++ b/proxypool/probe.go @@ -0,0 +1,90 @@ +package proxypool + +import ( + "context" + "fmt" + "io" + "net" + "net/http" + "net/url" + "strings" + "time" +) + +// DefaultProbeURL is the liveness target hit through each proxy. generate_204 is +// tiny, plain HTTP, and returns 204 with no body — cheap and unambiguous. +const DefaultProbeURL = "http://www.gstatic.com/generate_204" + +// ProbeResult is one liveness check outcome. +type ProbeResult struct { + OK bool + Latency time.Duration + Err string +} + +// Probe checks whether a proxy can reach probeURL within timeout, returning +// success + round-trip latency, or the failure reason. A 2xx/3xx response counts +// as alive (some probe targets redirect). +func Probe(ctx context.Context, proxyURL *url.URL, probeURL string, timeout time.Duration) ProbeResult { + if probeURL == "" { + probeURL = DefaultProbeURL + } + client, err := proxyHTTPClient(proxyURL, timeout) + if err != nil { + return ProbeResult{Err: err.Error()} + } + defer client.CloseIdleConnections() + + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, probeURL, nil) + if err != nil { + return ProbeResult{Err: err.Error()} + } + start := time.Now() + resp, err := client.Do(req) + if err != nil { + return ProbeResult{Err: trimErr(err.Error())} + } + defer resp.Body.Close() + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4096)) + latency := time.Since(start) + if resp.StatusCode >= 400 { + return ProbeResult{Err: fmt.Sprintf("probe status %d", resp.StatusCode), Latency: latency} + } + return ProbeResult{OK: true, Latency: latency} +} + +// proxyHTTPClient builds an *http.Client whose transport routes through proxyURL: +// http/https proxies use Transport.Proxy (native CONNECT), socks proxies use a +// custom DialContext. One-shot use — caller closes idle connections. +func proxyHTTPClient(proxyURL *url.URL, timeout time.Duration) (*http.Client, error) { + tr := &http.Transport{ + DisableKeepAlives: true, + TLSHandshakeTimeout: timeout, + } + switch strings.ToLower(proxyURL.Scheme) { + case "http", "https": + tr.Proxy = http.ProxyURL(proxyURL) + case "socks5": + tr.DialContext = func(ctx context.Context, _, address string) (net.Conn, error) { + return dialSOCKS(ctx, proxyURL, address) + } + default: + return nil, fmt.Errorf("unsupported proxy scheme %q", proxyURL.Scheme) + } + return &http.Client{Transport: tr, Timeout: timeout}, nil +} + +// trimErr shortens noisy dial errors to a storable single line. +func trimErr(s string) string { + s = strings.TrimSpace(s) + if i := strings.IndexByte(s, '\n'); i >= 0 { + s = s[:i] + } + const max = 200 + if len(s) > max { + s = s[:max] + } + return s +} diff --git a/proxypool/sources.go b/proxypool/sources.go new file mode 100644 index 0000000..8e53a3a --- /dev/null +++ b/proxypool/sources.go @@ -0,0 +1,104 @@ +package proxypool + +import ( + "bufio" + "bytes" + "context" + "fmt" + "io" + "net/http" + "strconv" + "strings" + "time" + + "github.com/Autumn-27/artex/db" +) + +// Source is one free proxy source. Only format-stable GitHub raw "host:port" lists +// are used (no HTML scraping) — each line is a bare host:port, protocol is fixed +// per source. Sources are DISABLED by default; users opt in per source. +type Source struct { + Name string + URL string + Protocol string // protocol assigned to every entry from this source +} + +// BuiltinSources is the catalog of free proxy sources. These are widely-mirrored, +// daily-updated raw text lists with a stable "ip:port\n" format. If a source dies +// (raw lists do rot), the user disables it; adding/removing entries here is the +// only maintenance touch-point. +// 仅收录可拨号的协议（http/https/socks5）——socks4 go-mitmproxy/x/net 均不支持，故不纳入。 +var BuiltinSources = []Source{ + {Name: "TheSpeedX-http", URL: "https://raw.githubusercontent.com/TheSpeedX/PROXY-List/master/http.txt", Protocol: "http"}, + {Name: "TheSpeedX-socks5", URL: "https://raw.githubusercontent.com/TheSpeedX/PROXY-List/master/socks5.txt", Protocol: "socks5"}, + {Name: "monosans-http", URL: "https://raw.githubusercontent.com/monosans/proxy-list/main/proxies/http.txt", Protocol: "http"}, + {Name: "monosans-socks5", URL: "https://raw.githubusercontent.com/monosans/proxy-list/main/proxies/socks5.txt", Protocol: "socks5"}, +} + +// SourceNames returns every built-in source name (for seeding the source list UI). +func SourceNames() []string { + out := make([]string, len(BuiltinSources)) + for i, s := range BuiltinSources { + out[i] = s.Name + } + return out +} + +// sourceByName looks up a built-in source definition. +func sourceByName(name string) (Source, bool) { + for _, s := range BuiltinSources { + if s.Name == name { + return s, true + } + } + return Source{}, false +} + +// fetch downloads and parses one source into proxy rows (connection fields only; +// health/quality are filled later by probing). +func fetch(ctx context.Context, src Source, client *http.Client) ([]db.Proxy, error) { + ctx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, src.URL, nil) + if err != nil { + return nil, err + } + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("source status %d", resp.StatusCode) + } + body, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20)) // 8MB cap + if err != nil { + return nil, err + } + return parseHostPortList(body, src.Protocol), nil +} + +// parseHostPortList turns a "host:port\n" list into proxy rows. Lines that are not +// a valid host:port are skipped. Protocol is fixed per source. +func parseHostPortList(body []byte, protocol string) []db.Proxy { + var out []db.Proxy + sc := bufio.NewScanner(bytes.NewReader(body)) + sc.Buffer(make([]byte, 0, 64*1024), 1<<20) + for sc.Scan() { + line := strings.TrimSpace(sc.Text()) + if line == "" || strings.HasPrefix(line, "#") { + continue + } + host, portStr, ok := strings.Cut(line, ":") + if !ok { + continue + } + host = strings.TrimSpace(host) + port, err := strconv.Atoi(strings.TrimSpace(portStr)) + if host == "" || err != nil || port <= 0 || port > 65535 { + continue + } + out = append(out, db.Proxy{Protocol: protocol, Host: host, Port: port}) + } + return out +} diff --git a/proxypool/sources_test.go b/proxypool/sources_test.go new file mode 100644 index 0000000..6e46771 --- /dev/null +++ b/proxypool/sources_test.go @@ -0,0 +1,43 @@ +package proxypool + +import "testing" + +// TestParseHostPortList verifies the raw "host:port" list parser skips blanks, +// comments, and malformed lines while tagging every entry with the source protocol. +func TestParseHostPortList(t *testing.T) { + body := []byte(` +1.2.3.4:8080 +# a comment +5.6.7.8:1080 + +9.9.9.9:notaport +10.0.0.1:70000 +11.22.33.44:3128 +`) + got := parseHostPortList(body, "socks5") + if len(got) != 3 { + t.Fatalf("want 3 valid proxies, got %d: %+v", len(got), got) + } + for _, p := range got { + if p.Protocol != "socks5" { + t.Errorf("protocol = %q, want socks5", p.Protocol) + } + } + if got[0].Host != "1.2.3.4" || got[0].Port != 8080 { + t.Errorf("first entry = %+v", got[0]) + } +} + +// TestSourceNamesUnique guards against a copy-paste duplicate in the catalog. +func TestSourceNamesUnique(t *testing.T) { + seen := map[string]bool{} + for _, n := range SourceNames() { + if seen[n] { + t.Fatalf("duplicate source name %q", n) + } + seen[n] = true + } + if len(seen) == 0 { + t.Fatal("no built-in sources defined") + } +} diff --git a/server/assembly.go b/server/assembly.go index 426a9b6..00daf37 100644 --- a/server/assembly.go +++ b/server/assembly.go @@ -342,6 +342,7 @@ func buildDomainReg(as *db.AssetStore) map[string]actool.CoreTool { } serverTS := agent.NewToolSet(nil, "") serverTS.SetAssetStore(as, as.Companies()) + serverTS.SetProxyStore(as.Proxies()) // list_proxies 若被 auto/自定义 agent 勾选，需要池句柄 reg := make(map[string]actool.CoreTool) for _, t := range serverTS.AllDomainTools() { reg[t.Name()] = t diff --git a/server/assets_scope_test.go b/server/assets_scope_test.go index c3e6b7c..f60944b 100644 --- a/server/assets_scope_test.go +++ b/server/assets_scope_test.go @@ -34,7 +34,7 @@ func TestCompanyScopeInputsAcceptStructuredAndLegacyRules(t *testing.T) { } func TestCreateCompanyRejectsNormalizedDuplicateWithoutChangingScope(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v)", err) } @@ -74,7 +74,7 @@ func TestCreateCompanyRejectsNormalizedDuplicateWithoutChangingScope(t *testing. } func TestDeleteCompanyRefreshesLiveTaskCompanyIDs(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v)", err) } @@ -115,7 +115,7 @@ func TestDeleteCompanyRefreshesLiveTaskCompanyIDs(t *testing.T) { } func TestCompanyScopeHTTPErrorClassificationAndBounds(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v)", err) } @@ -180,7 +180,7 @@ func TestCompanyScopeHTTPErrorClassificationAndBounds(t *testing.T) { } func TestListAssetsClassifiesValidationAndDatabaseErrors(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v)", err) } @@ -206,7 +206,7 @@ func TestListAssetsClassifiesValidationAndDatabaseErrors(t *testing.T) { } func TestDeleteCompanyRejectsBadJSONAndReportsMissing(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v)", err) } @@ -245,7 +245,7 @@ func TestDeleteCompanyRejectsBadJSONAndReportsMissing(t *testing.T) { } func TestCompanyScopeSystemFailureIsHTTP500(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v)", err) } diff --git a/server/core_test.go b/server/core_test.go index 56b8a18..6754518 100644 --- a/server/core_test.go +++ b/server/core_test.go @@ -18,7 +18,7 @@ import ( // TestCoreTaskLifecyclePG exercises the migrated core (tasks/exploration on PG) // through the real HTTP mux: create → goal nodes seeded → list → delete cascade. func TestCoreTaskLifecyclePG(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) — skipping", err) } diff --git a/server/findings_groups_test.go b/server/findings_groups_test.go index 5cb44ed..4af99ae 100644 --- a/server/findings_groups_test.go +++ b/server/findings_groups_test.go @@ -37,7 +37,7 @@ func TestFindingPaginationParam(t *testing.T) { } func TestFindingGroupsReturnsTaskBucketsAndNormalizesPagination(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) — skipping", err) } @@ -125,7 +125,7 @@ func TestFindingGroupsReturnsTaskBucketsAndNormalizesPagination(t *testing.T) { } func TestDeepenFindingCreatesAuditedIntentAndRevivesTask(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) — skipping", err) } @@ -254,7 +254,7 @@ func TestDeepenFindingValidatesDescription(t *testing.T) { } func TestDeepenAdmissionFailureDiscardsFollowUpIntent(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } diff --git a/server/inheritance_api_test.go b/server/inheritance_api_test.go index df9fe4f..807207c 100644 --- a/server/inheritance_api_test.go +++ b/server/inheritance_api_test.go @@ -12,7 +12,7 @@ import ( ) func TestInheritedActivityDetailAndRelationDeletion(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } diff --git a/server/manager.go b/server/manager.go index 4bac104..ed35d4d 100644 --- a/server/manager.go +++ b/server/manager.go @@ -1,11 +1,13 @@ package server import ( + "context" "database/sql" "encoding/json" "errors" "fmt" "log" + "net/url" "os" "path/filepath" "sort" @@ -19,6 +21,7 @@ import ( "github.com/Autumn-27/artex/enrich" "github.com/Autumn-27/artex/guard" "github.com/Autumn-27/artex/intercept" + "github.com/Autumn-27/artex/proxypool" "github.com/Autumn-27/artex/traffic" actool "github.com/Autumn-27/norma/tool" ) @@ -224,6 +227,9 @@ type Manager struct { braveKey string tavilyKey string webSearchProxy string + proxyPool *proxypool.Pool // 出口代理池后台任务（抓取 + 验活） + proxyGW *proxypool.Gateway // 入口C：本地转发网关（纯 TCP，不解密） + proxyGWAddr string // 网关监听地址（空 = 禁用入口C） } // Settings keys the UI toggles at runtime. @@ -245,8 +251,17 @@ const ( // 任务并发上限:开关 + 上限数。默认关闭;开启后默认上限 5(见 defaultConcurrencyLimit)。 settingConcurrencyOn = "task_concurrency_enabled" settingConcurrencyLimit = "task_concurrency_limit" + // 代理池：主开关、"主出口只走可信代理"安全阀门、抓取/验活间隔（分钟）。 + // 主开关键复用 db.SettingProxyPoolEnabled，agent 工具层与此读同一个键。 + settingProxyPoolOn = pgdb.SettingProxyPoolEnabled + settingProxyTrustedOnly = "proxy_egress_trusted_only" + settingProxyFetchMin = "proxy_fetch_interval_min" + settingProxyCheckMin = "proxy_check_interval_min" // defaultWebSearchBackend is used when web search is on but no backend was picked. defaultWebSearchBackend = "ddgs" + // 代理池默认间隔：免费代理失效快，验活默认 30min、抓取默认 15min。 + defaultProxyFetchMin = 15 + defaultProxyCheckMin = 30 // defaultWorkers is the concurrent work-agent count when the setting is unset. defaultWorkers = 3 // defaultConcurrencyLimit is the simultaneous-running-task cap when the feature @@ -308,7 +323,7 @@ func (m *Manager) Enrich() *enrich.Engine { return m.enrich } // NewManager connects to PostgreSQL and, if proxyAddr is non-empty, starts the // traffic-recording proxy. PostgreSQL is required (it is the single data source). -func NewManager(dir, proxyAddr string) (*Manager, error) { +func NewManager(dir, proxyAddr, gatewayAddr string) (*Manager, error) { // Resolve the data dir to an ABSOLUTE path up front. Every data path derives // from it — notably the MITM CA cert, whose path is injected into worker shells // (SSL_CERT_FILE/CURL_CA_BUNDLE) and read by WebFetch. A relative path (the @@ -336,13 +351,16 @@ func NewManager(dir, proxyAddr string) (*Manager, error) { if err := pg.EnsureLLMUsageTable(); err != nil { log.Printf("[llmusage] create table: %v", err) } - m := &Manager{dir: dir, pg: pg, assets: pg.Assets(), tasks: map[string]*Task{}, interceptor: intercept.New(pg)} + m := &Manager{dir: dir, pg: pg, assets: pg.Assets(), tasks: map[string]*Task{}, interceptor: intercept.New(pg), proxyGWAddr: gatewayAddr} if proxyAddr != "" { tr, err := traffic.Open(filepath.Join(dir, "traffic"), proxyAddr) if err != nil { log.Printf("[traffic] disabled: %v", err) } else { m.traffic = tr + // 入口A：代理池开启时，MITM 上游按目标 host 从池里选一个出口（每 host 粘性， + // 受"仅可信"安全阀门约束）；池关或无可用代理时返回 nil → 直连目标。 + tr.SetPoolUpstream(m.selectPoolUpstream) go func() { log.Printf("[traffic] recording proxy on %s (set HTTP_PROXY=%s + trust _ca CA)", proxyAddr, tr.ProxyAddr()) if err := tr.Start(); err != nil { @@ -376,9 +394,42 @@ func NewManager(dir, proxyAddr string) (*Manager, error) { // Reconcile the seeded browser MCP with the persisted capture state, so a // restart with capture already on keeps Playwright routed through the proxy. m.syncBrowserMCPProxy() + // Outbound proxy pool background loops (fetch free sources + probe liveness). + // Loops read the master switch each tick; they no-op while the pool is off. + m.proxyPool = proxypool.NewPool(pg, proxypool.Config{ + Enabled: m.ProxyPoolEnabled, + FetchInterval: func() time.Duration { return time.Duration(m.proxyFetchMin()) * time.Minute }, + CheckInterval: func() time.Duration { return time.Duration(m.proxyCheckMin()) * time.Minute }, + }) + m.proxyPool.Start(context.Background()) + // 入口C：本地转发网关（纯 TCP，不解密、不抓包）。常驻监听；仅当 MITM 关 + 池开时 + // ProxyAddr() 才把它注入给 agent（见下）。上游选择与入口A 共用 selectPoolUpstream。 + if gatewayAddr != "" { + gw := proxypool.NewGateway(gatewayAddr) + gw.SetUpstream(m.selectPoolUpstream) + if err := gw.Start(); err != nil { + log.Printf("[proxypool] gateway disabled: %v", err) + } else { + m.proxyGW = gw + } + } return m, nil } +// selectPoolUpstream picks a pool proxy URL for a target host, or nil to dial +// direct. Shared by the MITM upstream (入口A) and the gateway (入口C): honors the +// pool master switch and the trusted-only egress safety valve. +func (m *Manager) selectPoolUpstream(host string) *url.URL { + if !m.ProxyPoolEnabled() { + return nil + } + p, err := m.pg.Proxies().SelectForHost(host, m.ProxyEgressTrustedOnly()) + if err != nil || p == nil { + return nil + } + return p.URL() +} + // TrafficEnabled reports whether traffic capture is on (default off). When off, // no proxy/traffic tools/prompt are injected into agents (nothing is recorded). func (m *Manager) TrafficEnabled() bool { @@ -638,27 +689,104 @@ func (m *Manager) Assets() *pgdb.AssetStore { return m.assets } func (m *Manager) PG() *pgdb.DB { return m.pg } func (m *Manager) Traffic() *traffic.Traffic { return m.traffic } -// ProxyAddr returns the recording proxy address agents route through — empty when -// traffic capture is off, so no proxy is injected (agent runs direct, no recording). +// Proxies returns the outbound proxy pool store. +func (m *Manager) Proxies() *pgdb.ProxyStore { return m.pg.Proxies() } + +// ProxyPool returns the background pool runner (fetch + probe + manual check). +func (m *Manager) ProxyPool() *proxypool.Pool { return m.proxyPool } + +// ProxyPoolEnabled reports whether the outbound proxy pool is on (default off). +func (m *Manager) ProxyPoolEnabled() bool { return m.pg.GetBool(settingProxyPoolOn, false) } + +// SetProxyPoolEnabled toggles the pool master switch. +func (m *Manager) SetProxyPoolEnabled(on bool) error { return m.pg.SetBool(settingProxyPoolOn, on) } + +// ProxyEgressTrustedOnly reports whether the main egress (MITM upstream / gateway) +// only rotates trusted (manual/imported) proxies. Default true — free-pool nodes +// stay out of the main egress unless the user opts them in. +func (m *Manager) ProxyEgressTrustedOnly() bool { + return m.pg.GetBool(settingProxyTrustedOnly, true) +} + +// SetProxyEgressTrustedOnly toggles the trusted-only egress safety valve. +func (m *Manager) SetProxyEgressTrustedOnly(on bool) error { + return m.pg.SetBool(settingProxyTrustedOnly, on) +} + +// proxyFetchMin / proxyCheckMin read the pool intervals (minutes), falling back to +// defaults when unset or non-positive. +func (m *Manager) proxyFetchMin() int { return m.settingIntDefault(settingProxyFetchMin, defaultProxyFetchMin) } +func (m *Manager) proxyCheckMin() int { return m.settingIntDefault(settingProxyCheckMin, defaultProxyCheckMin) } + +// ProxyIntervals returns the fetch/check intervals in minutes (for the settings UI). +func (m *Manager) ProxyIntervals() (fetchMin, checkMin int) { + return m.proxyFetchMin(), m.proxyCheckMin() +} + +// SetProxyIntervals persists the fetch/check intervals (minutes); non-positive +// values are ignored so a partial update keeps the other. +func (m *Manager) SetProxyIntervals(fetchMin, checkMin int) error { + if fetchMin > 0 { + if err := m.pg.SetSetting(settingProxyFetchMin, strconv.Itoa(fetchMin)); err != nil { + return err + } + } + if checkMin > 0 { + if err := m.pg.SetSetting(settingProxyCheckMin, strconv.Itoa(checkMin)); err != nil { + return err + } + } + return nil +} + +// settingIntDefault reads an integer setting, returning def when unset/invalid. +func (m *Manager) settingIntDefault(key string, def int) int { + v, ok, err := m.pg.GetSetting(key) + if err != nil || !ok || v == "" { + return def + } + n, err := strconv.Atoi(v) + if err != nil || n <= 0 { + return def + } + return n +} + +// ProxyAddr returns the local proxy address agents route through, decided by the +// MITM(traffic) × proxy-pool switch matrix: +// - MITM on → the recording proxy (入口A; pool-on makes its upstream rotate) +// - MITM off + pool on → the pool gateway (入口C; plain forward, no capture) +// - both off → "" (direct, nothing injected) func (m *Manager) ProxyAddr() string { - if m.traffic == nil || !m.TrafficEnabled() { - return "" + if m.traffic != nil && m.TrafficEnabled() { + return m.traffic.ProxyAddr() + } + if m.proxyGW != nil && m.ProxyPoolEnabled() { + return "http://" + m.proxyGW.Addr() } - return m.traffic.ProxyAddr() + return "" } -// ProxyCACert returns the recording proxy's CA cert path (empty when no proxy or -// traffic capture is off), which WebFetch trusts to verify HTTPS through the MITM. +// ProxyCACert returns the CA the injected proxy needs trusted. Only the MITM +// recording proxy terminates TLS and needs its CA; the gateway is a plain tunnel +// that never decrypts, so it MUST inject no CA (the agent does real TLS to the +// target). Empty when direct or in gateway mode. func (m *Manager) ProxyCACert() string { - if m.traffic == nil || !m.TrafficEnabled() { - return "" + if m.traffic != nil && m.TrafficEnabled() { + return m.traffic.CACertPath() } - return m.traffic.CACertPath() + return "" } func (m *Manager) Close() error { m.mu.Lock() defer m.mu.Unlock() + if m.proxyPool != nil { + m.proxyPool.Stop() + } + if m.proxyGW != nil { + m.proxyGW.Stop() + } if m.traffic != nil { m.traffic.Close() } diff --git a/server/manager_delete_files_test.go b/server/manager_delete_files_test.go index 48de983..becaaa2 100644 --- a/server/manager_delete_files_test.go +++ b/server/manager_delete_files_test.go @@ -13,7 +13,7 @@ import ( ) func TestSeedAssociatesTargetAssetWithTask(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -172,7 +172,7 @@ func TestStageTaskFilesRollbackReportsRestoreFailure(t *testing.T) { func TestManagerDeleteTaskRestoresFilesAndTrafficWhenPostgresDeleteFails(t *testing.T) { dataDir := t.TempDir() - m, err := NewManager(dataDir, "") + m, err := NewManager(dataDir, "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } diff --git a/server/mgmt_test.go b/server/mgmt_test.go index 9bfebb7..db4525c 100644 --- a/server/mgmt_test.go +++ b/server/mgmt_test.go @@ -11,7 +11,7 @@ import ( // TestMgmtAPI exercises the PostgreSQL-backed management API through the real mux. func TestMgmtAPI(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("database unavailable (%v) — skipping management API test", err) } diff --git a/server/orchestration.go b/server/orchestration.go index 5ee78f0..6c65cd4 100644 --- a/server/orchestration.go +++ b/server/orchestration.go @@ -162,6 +162,7 @@ func (s *Server) delegateToTask(ctx context.Context, in json.RawMessage, pick fu tsx := agent.NewToolSet(t.Store, "orchestrator") if s.m.Assets() != nil { tsx.SetAssetStore(s.m.Assets(), s.m.Assets().Companies()) + tsx.SetProxyStore(s.m.Assets().Proxies()) } tsx.SetNotify(t.Notify) // hint writes wake this task's planner (no-op for read tools) return pick(tsx).Call(ctx, inner, nil) diff --git a/server/proxies.go b/server/proxies.go new file mode 100644 index 0000000..7c2ae1f --- /dev/null +++ b/server/proxies.go @@ -0,0 +1,228 @@ +package server + +import ( + "encoding/json" + "errors" + "net/http" + "strconv" + "strings" + + "github.com/Autumn-27/artex/db" + "github.com/Autumn-27/artex/proxypool" +) + +const maxProxyBodyBytes = 4 << 20 + +// maskProxies returns copies safe to serialize (password replaced with a +// placeholder), so proxy credentials never leave the backend. +func maskProxies(in []*db.Proxy) []db.Proxy { + out := make([]db.Proxy, len(in)) + for i, p := range in { + out[i] = p.Masked() + } + return out +} + +// listProxies GET /api/proxies — filters ?protocol=®ion=&tag=&healthy=1&enabled=1 +// plus server-side paging ?page=&limit= (page 1-based; omit for all rows). +func (s *Server) listProxies(w http.ResponseWriter, r *http.Request) { + q := r.URL.Query() + f := db.ProxyFilter{ + Protocol: q.Get("protocol"), + Region: q.Get("region"), + Anonymity: q.Get("anonymity"), + OnlyHealthy: q.Get("healthy") == "1", + OnlyEnabled: q.Get("enabled") == "1", + } + if tag := q.Get("tag"); tag != "" { + f.Tags = []string{tag} + } + page, _ := strconv.Atoi(q.Get("page")) + limit, _ := strconv.Atoi(q.Get("limit")) + if limit > 0 { + if page < 1 { + page = 1 + } + f.Limit, f.Offset = limit, (page-1)*limit + } + store := s.m.Proxies() + total, err := store.CountProxies(f) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + rows, err := store.ListProxies(f) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + writeJSON(w, 200, map[string]any{"proxies": maskProxies(rows), "total": total}) +} + +// proxyCreateReq is the create/update body. Password "" on update keeps the +// existing secret (so the masked value returned by GET is never written back). +type proxyCreateReq struct { + Protocol string `json:"protocol"` + Host string `json:"host"` + Port int `json:"port"` + Username string `json:"username"` + Password string `json:"password"` + Anonymity string `json:"anonymity"` + Region string `json:"region"` + Tags []string `json:"tags"` + Label string `json:"label"` + Enabled bool `json:"enabled"` +} + +func (s *Server) decodeProxyBody(w http.ResponseWriter, r *http.Request, v any) bool { + r.Body = http.MaxBytesReader(w, r.Body, maxProxyBodyBytes) + if err := json.NewDecoder(r.Body).Decode(v); err != nil { + writeErr(w, http.StatusBadRequest, "invalid JSON: "+err.Error()) + return false + } + return true +} + +// createProxy POST /api/proxies — add one manual proxy (trusted). +func (s *Server) createProxy(w http.ResponseWriter, r *http.Request) { + var req proxyCreateReq + if !s.decodeProxyBody(w, r, &req) { + return + } + p := &db.Proxy{ + Protocol: req.Protocol, Host: strings.TrimSpace(req.Host), Port: req.Port, + Username: req.Username, Password: req.Password, Anonymity: req.Anonymity, + Region: req.Region, Tags: req.Tags, Label: req.Label, + Source: "manual", Trusted: true, Enabled: req.Enabled, + } + id, err := s.m.Proxies().CreateProxy(p) + if err != nil { + writeErr(w, 400, err.Error()) + return + } + writeJSON(w, 200, map[string]any{"id": id}) +} + +// updateProxy PUT /api/proxies/{id} — edit one proxy. A blank password preserves +// the stored secret (the UI shows a masked value it must not persist back). +func (s *Server) updateProxy(w http.ResponseWriter, r *http.Request) { + id, err := strconv.ParseInt(r.PathValue("id"), 10, 64) + if err != nil { + writeErr(w, 400, "invalid id") + return + } + var req proxyCreateReq + if !s.decodeProxyBody(w, r, &req) { + return + } + cur, err := s.m.Proxies().GetProxy(id) + if err != nil { + writeProxyErr(w, err) + return + } + cur.Protocol, cur.Host, cur.Port = req.Protocol, strings.TrimSpace(req.Host), req.Port + cur.Username, cur.Anonymity, cur.Region = req.Username, req.Anonymity, req.Region + cur.Tags, cur.Label, cur.Enabled = req.Tags, req.Label, req.Enabled + if req.Password != "" && req.Password != "********" { + cur.Password = req.Password + } + if err := s.m.Proxies().UpdateProxy(cur); err != nil { + writeProxyErr(w, err) + return + } + writeJSON(w, 200, map[string]any{"ok": true}) +} + +// deleteProxy DELETE /api/proxies/{id}. +func (s *Server) deleteProxy(w http.ResponseWriter, r *http.Request) { + id, err := strconv.ParseInt(r.PathValue("id"), 10, 64) + if err != nil { + writeErr(w, 400, "invalid id") + return + } + if err := s.m.Proxies().DeleteProxy(id); err != nil { + writeProxyErr(w, err) + return + } + writeJSON(w, 200, map[string]any{"ok": true}) +} + +// importProxies POST /api/proxies/import — body {"text":"host:port\n..."}; adds +// each parsed line as a trusted import, de-duped by (protocol,host,port). +func (s *Server) importProxies(w http.ResponseWriter, r *http.Request) { + var req struct { + Text string `json:"text"` + } + if !s.decodeProxyBody(w, r, &req) { + return + } + lines := strings.Split(req.Text, "\n") + added, invalid, err := s.m.Proxies().ImportProxies(lines) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + writeJSON(w, 200, map[string]any{"added": added, "invalid": invalid}) +} + +// checkProxy POST /api/proxies/{id}/check — probe one proxy on demand. +func (s *Server) checkProxy(w http.ResponseWriter, r *http.Request) { + id, err := strconv.ParseInt(r.PathValue("id"), 10, 64) + if err != nil { + writeErr(w, 400, "invalid id") + return + } + res, err := s.m.ProxyPool().ProbeNow(r.Context(), id) + if err != nil { + writeProxyErr(w, err) + return + } + writeJSON(w, 200, map[string]any{ + "ok": res.OK, "latency_ms": res.Latency.Milliseconds(), "error": res.Err, + }) +} + +// listProxySources GET /api/proxy-sources — built-in source catalog + enable state. +func (s *Server) listProxySources(w http.ResponseWriter, r *http.Request) { + sources, err := s.m.Proxies().ListSources(proxypool.SourceNames()) + if err != nil { + writeErr(w, 500, err.Error()) + return + } + writeJSON(w, 200, map[string]any{"sources": sources}) +} + +// fetchProxySource POST /api/proxy-sources/{name}/fetch — pull one free source now. +func (s *Server) fetchProxySource(w http.ResponseWriter, r *http.Request) { + name := r.PathValue("name") + total, added, err := s.m.ProxyPool().FetchSourceNow(r.Context(), name) + if err != nil { + writeErr(w, 400, err.Error()) + return + } + writeJSON(w, 200, map[string]any{"fetched": total, "added": added}) +} + +// setProxySource PUT /api/proxy-sources/{name} — body {"enabled":bool}. +func (s *Server) setProxySource(w http.ResponseWriter, r *http.Request) { + name := r.PathValue("name") + var req struct { + Enabled bool `json:"enabled"` + } + if !s.decodeProxyBody(w, r, &req) { + return + } + if err := s.m.Proxies().SetSourceEnabled(name, req.Enabled); err != nil { + writeErr(w, 500, err.Error()) + return + } + writeJSON(w, 200, map[string]any{"ok": true}) +} + +func writeProxyErr(w http.ResponseWriter, err error) { + if errors.Is(err, db.ErrProxyNotFound) { + writeErr(w, 404, "proxy not found") + return + } + writeErr(w, 500, err.Error()) +} diff --git a/server/server.go b/server/server.go index f376df2..a004d83 100644 --- a/server/server.go +++ b/server/server.go @@ -688,6 +688,16 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("DELETE /api/companies/{id}", s.deleteCompany) mux.HandleFunc("POST /api/companies/{id}/scope", s.addCompanyScope) mux.HandleFunc("POST /api/companies/reattribute", s.reattribute) + // 代理池：条目 CRUD/导入/探活 + 免费源开关 + mux.HandleFunc("GET /api/proxies", s.listProxies) + mux.HandleFunc("POST /api/proxies", s.createProxy) + mux.HandleFunc("POST /api/proxies/import", s.importProxies) + mux.HandleFunc("PUT /api/proxies/{id}", s.updateProxy) + mux.HandleFunc("DELETE /api/proxies/{id}", s.deleteProxy) + mux.HandleFunc("POST /api/proxies/{id}/check", s.checkProxy) + mux.HandleFunc("GET /api/proxy-sources", s.listProxySources) + mux.HandleFunc("PUT /api/proxy-sources/{name}", s.setProxySource) + mux.HandleFunc("POST /api/proxy-sources/{name}/fetch", s.fetchProxySource) mux.HandleFunc("GET /api/exploration/frontier", s.frontier) mux.HandleFunc("GET /api/exploration/findings", s.findings) @@ -2978,6 +2988,7 @@ func (s *Server) settingsPayload() map[string]any { if concLimit == 0 { concLimit = defaultConcurrencyLimit // 关闭时也回显一个合理默认值给 UI } + fetchMin, checkMin := s.m.ProxyIntervals() return map[string]any{ "traffic_capture": s.m.TrafficEnabled(), "llm_record": s.m.LLMRecordEnabled(), @@ -2994,6 +3005,11 @@ func (s *Server) settingsPayload() map[string]any { // 自动切到下一个配置。bind_fallback 仅在轮询开启时有意义(默认关)。 "llm_pool_enabled": s.m.LLMPoolEnabled(), "llm_pool_bind_fallback": s.m.LLMPoolBindFallback(), + // 代理池：主开关；"主出口只走可信代理"安全阀门；抓取/验活间隔(分钟)。 + "proxy_pool_enabled": s.m.ProxyPoolEnabled(), + "proxy_egress_trusted_only": s.m.ProxyEgressTrustedOnly(), + "proxy_fetch_interval_min": fetchMin, + "proxy_check_interval_min": checkMin, } } @@ -3034,6 +3050,12 @@ func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) { // 重建 provider 链才生效，走下面的 changed → applyLLM 路径。 LLMPoolEnabled *bool `json:"llm_pool_enabled"` LLMPoolBindFallback *bool `json:"llm_pool_bind_fallback"` + // 代理池：主开关、"仅可信"安全阀门、抓取/验活间隔(分钟)。即时生效，无需重建 agent + // （出口注入在每次取 ProxyAddr 时按开关计算；后台 loop 每 tick 读间隔）。 + ProxyPoolEnabled *bool `json:"proxy_pool_enabled"` + ProxyEgressTrustedOnly *bool `json:"proxy_egress_trusted_only"` + ProxyFetchIntervalMin *int `json:"proxy_fetch_interval_min"` + ProxyCheckIntervalMin *int `json:"proxy_check_interval_min"` } if err := json.NewDecoder(r.Body).Decode(&req); err != nil { writeErr(w, 400, err.Error()) @@ -3105,6 +3127,34 @@ func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) { } changed = true } + if req.ProxyPoolEnabled != nil { + if err := s.m.SetProxyPoolEnabled(*req.ProxyPoolEnabled); err != nil { + writeErr(w, 500, err.Error()) + return + } + // 开关改变 ProxyAddr() 注入的出口地址（traffic 关时：网关 vs 直连），需重建 + // agent 让新出口生效——与 traffic_capture 同理。 + changed = true + } + if req.ProxyEgressTrustedOnly != nil { + if err := s.m.SetProxyEgressTrustedOnly(*req.ProxyEgressTrustedOnly); err != nil { + writeErr(w, 500, err.Error()) + return + } + } + if req.ProxyFetchIntervalMin != nil || req.ProxyCheckIntervalMin != nil { + fetchMin, checkMin := 0, 0 + if req.ProxyFetchIntervalMin != nil { + fetchMin = *req.ProxyFetchIntervalMin + } + if req.ProxyCheckIntervalMin != nil { + checkMin = *req.ProxyCheckIntervalMin + } + if err := s.m.SetProxyIntervals(fetchMin, checkMin); err != nil { + writeErr(w, 500, err.Error()) + return + } + } if req.WebSearchEnabled != nil || req.WebSearchBackend != nil || req.BraveKey != nil || req.TavilyKey != nil || req.WebSearchProxy != nil { // Fill unspecified fields from current state so a partial PUT doesn't reset them. on, backend, _, _, _ := s.m.WebSearch() diff --git a/server/task_admission_test.go b/server/task_admission_test.go index 6a3e1ef..457accf 100644 --- a/server/task_admission_test.go +++ b/server/task_admission_test.go @@ -43,7 +43,7 @@ func restoreConcurrencySetting(t *testing.T, m *Manager) func() { } func TestAdmitAlreadyRunningTaskDoesNotQueueAfterLimitDecrease(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -85,7 +85,7 @@ func TestAdmitAlreadyRunningTaskDoesNotQueueAfterLimitDecrease(t *testing.T) { } func TestAdmissionDoesNotOverwriteConcurrentTerminalStatus(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -122,7 +122,7 @@ func TestAdmissionDoesNotOverwriteConcurrentTerminalStatus(t *testing.T) { } func TestAdmitKeepsQueuedBootstrapMode(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -165,7 +165,7 @@ func TestAdmitKeepsQueuedBootstrapMode(t *testing.T) { } func TestTerminalTaskQueuedByAdmissionKeepsExecutionBarrier(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -213,7 +213,7 @@ func TestTerminalTaskQueuedByAdmissionKeepsExecutionBarrier(t *testing.T) { } func TestTimedOutTaskRevivalResetsClockAndSettlingAcrossFIFO(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -299,7 +299,7 @@ func TestTimedOutTaskRevivalResetsClockAndSettlingAcrossFIFO(t *testing.T) { } func TestQueuedTaskCanPauseAndResumeAtFIFOTail(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -349,7 +349,7 @@ func TestQueuedTaskCanPauseAndResumeAtFIFOTail(t *testing.T) { } func TestReadyFIFOIsAdmittedBeforeNewTask(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -396,7 +396,7 @@ func TestReadyFIFOIsAdmittedBeforeNewTask(t *testing.T) { } func TestUnavailableTaskReleasesSlotForReadyQueue(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -435,7 +435,7 @@ func TestUnavailableTaskReleasesSlotForReadyQueue(t *testing.T) { } func TestUnavailableTaskWaitsForActiveLLMCallBeforeParking(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -489,7 +489,7 @@ func TestUnavailableTaskWaitsForActiveLLMCallBeforeParking(t *testing.T) { } func TestRerunAdmissionFailureRestoresIntentState(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -533,7 +533,7 @@ func TestRerunAdmissionFailureRestoresIntentState(t *testing.T) { } func TestTaskLLMResolutionRejectsInvalidExplicitProfile(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -574,7 +574,7 @@ func TestTaskLLMResolutionRejectsInvalidExplicitProfile(t *testing.T) { // An Agent binding outranks the task's LLM chain; roles left unbound keep using // the chain, so binding one role does not move the others off it. func TestTaskLLMResolutionPrefersAgentBindingOverTaskChain(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -641,7 +641,7 @@ func TestTaskLLMResolutionPrefersAgentBindingOverTaskChain(t *testing.T) { } func TestTaskLLMResolutionReportsDatabaseFailure(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -686,7 +686,7 @@ func TestTokenStatsWithoutTaskReturnsStableShape(t *testing.T) { } func TestBatchControlLimitCountsDeduplicatedIDs(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -760,7 +760,7 @@ func TestBatchControlLimitCountsDeduplicatedIDs(t *testing.T) { } func TestPauseTaskToolPersistsAndDequeuesTask(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } @@ -796,7 +796,7 @@ func TestPauseTaskToolPersistsAndDequeuesTask(t *testing.T) { } func TestPauseTaskToolUsesOrchestratorCancellationCause(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } diff --git a/server/task_control_routes_test.go b/server/task_control_routes_test.go index 26e2ad6..13d9985 100644 --- a/server/task_control_routes_test.go +++ b/server/task_control_routes_test.go @@ -9,7 +9,7 @@ import ( ) func TestWorkerControlRoutes(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } diff --git a/server/task_delete_barrier_test.go b/server/task_delete_barrier_test.go index 361b161..eaf4d45 100644 --- a/server/task_delete_barrier_test.go +++ b/server/task_delete_barrier_test.go @@ -145,7 +145,7 @@ func TestTaskLifecycleRechecksDeleteBarrierAfterConcLock(t *testing.T) { } func TestAbortTaskDeleteUsesPersistedPauseAndQueueState(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } diff --git a/server/task_llm_test.go b/server/task_llm_test.go index 73f28f2..e2d4d74 100644 --- a/server/task_llm_test.go +++ b/server/task_llm_test.go @@ -476,7 +476,7 @@ func TestTaskLLMStreamStopsWhenChainExhausted(t *testing.T) { } func TestProfileDeleteRestoresQuotaBlockedIntentWhenFallbackAvailable(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) - skipping", err) } diff --git a/server/task_templates_test.go b/server/task_templates_test.go index a3a9b1b..ab810a5 100644 --- a/server/task_templates_test.go +++ b/server/task_templates_test.go @@ -15,7 +15,7 @@ import ( ) func TestTaskTemplateHTTPCRUD(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) — skipping", err) } @@ -85,7 +85,7 @@ func TestTaskTemplateHTTPCRUD(t *testing.T) { } func TestConversationPatchReturnsPinState(t *testing.T) { - m, err := NewManager(t.TempDir(), "") + m, err := NewManager(t.TempDir(), "", "") if err != nil { t.Skipf("postgres unavailable (%v) — skipping", err) } diff --git a/traffic/traffic.go b/traffic/traffic.go index e827443..6ea89d0 100644 --- a/traffic/traffic.go +++ b/traffic/traffic.go @@ -120,6 +120,21 @@ type Traffic struct { // reason; connections to them are tunneled transparently (fail-open) so the // request still reaches the target — unrecorded — instead of being killed. pass sync.Map // hostname(string) -> struct{} + // poolUpstream, if set, picks an upstream proxy for a target host so recorded + // traffic egresses through the proxy pool (入口A). nil / returns nil = dial the + // target directly. Set by the server; the pool switch/selection lives there. + poolUpstream atomic.Pointer[func(host string) *url.URL] +} + +// SetPoolUpstream installs (or clears with nil) the proxy-pool upstream resolver. +// When set and it returns a non-nil URL for a host, recorded target traffic to +// that host is forwarded through the returned proxy instead of dialed directly. +func (t *Traffic) SetPoolUpstream(fn func(host string) *url.URL) { + if fn == nil { + t.poolUpstream.Store(nil) + return + } + t.poolUpstream.Store(&fn) } // Open initializes the traffic tree, blob store and SQLite index under dir. @@ -159,12 +174,19 @@ func Open(dir, addr string) (*Traffic, error) { db.Close() return nil, err } - // Dial targets DIRECTLY. go-mitmproxy's default upstream uses - // http.ProxyFromEnvironment, so an HTTP_PROXY/HTTPS_PROXY in the environment - // (a VPN/system proxy) would make it forward target requests through that - // external proxy — which can't reach the target → 502. We capture target - // traffic directly, never via the host's proxy. - p.SetUpstreamProxy(func(*http.Request) (*url.URL, error) { return nil, nil }) + // Upstream selection: by default dial targets DIRECTLY (never via the host's + // HTTP_PROXY/HTTPS_PROXY — go-mitmproxy's default would forward through a + // VPN/system proxy that can't reach the target → 502). When the proxy pool is + // on (入口A), poolUpstream returns a pool proxy for the host so recorded traffic + // egresses through it; a nil return falls back to direct. + p.SetUpstreamProxy(func(req *http.Request) (*url.URL, error) { + if fn := t.poolUpstream.Load(); fn != nil { + if u := (*fn)(hostOnly(req.Host)); u != nil { + return u, nil + } + } + return nil, nil + }) // Fail-open: MITM every host by default, EXCEPT ones a prior request proved we // can't intercept without breaking (see maybePassthrough). Those are tunneled // transparently so the request still reaches the target instead of being killed. diff --git a/web/src/app/(main)/system/proxies/page.tsx b/web/src/app/(main)/system/proxies/page.tsx new file mode 100644 index 0000000..e4aebbe --- /dev/null +++ b/web/src/app/(main)/system/proxies/page.tsx @@ -0,0 +1,516 @@ +"use client"; + +import * as React from "react"; + +import { + CircleCheckIcon, + CircleXIcon, + DownloadIcon, + Loader2Icon, + PencilIcon, + PlusIcon, + RefreshCwIcon, + Trash2Icon, +} from "lucide-react"; +import { toast } from "sonner"; + +import { TablePagination } from "@/components/table-pagination"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { NativeSelect, NativeSelectOption } from "@/components/ui/native-select"; +import { Switch } from "@/components/ui/switch"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { Textarea } from "@/components/ui/textarea"; +import { api } from "@/lib/api"; +import type { ProxyInput, ProxyNode, ProxyProtocol, ProxySource } from "@/lib/types"; + +const PROTOCOLS: ProxyProtocol[] = ["http", "https", "socks5"]; + +function emptyInput(): ProxyInput { + return { protocol: "http", host: "", port: 8080, enabled: true, tags: [] }; +} + +// 健康徽标：健康显示延迟，不健康显示原因，未检测显示灰点。 +function HealthBadge({ p }: { p: ProxyNode }) { + if (p.check_count === 0) { + return 未检测; + } + if (p.healthy) { + return ( + + {p.latency_ms}ms + + ); + } + return ( + + 失败 + + ); +} + +export default function ProxyPoolPage() { + const [proxies, setProxies] = React.useState([]); + const [total, setTotal] = React.useState(0); + const [page, setPage] = React.useState(1); + const [pageSize, setPageSize] = React.useState(20); + const [sources, setSources] = React.useState([]); + const [loading, setLoading] = React.useState(true); + const [checking, setChecking] = React.useState>(new Set()); + const [fetching, setFetching] = React.useState>(new Set()); + const [editing, setEditing] = React.useState(null); + const [form, setForm] = React.useState(null); + const [importOpen, setImportOpen] = React.useState(false); + const [importText, setImportText] = React.useState(""); + const [saving, setSaving] = React.useState(false); + + const load = React.useCallback(async () => { + try { + const [px, srcs] = await Promise.all([api.proxies({ page, limit: pageSize }), api.proxySources()]); + setProxies(px.proxies); + setTotal(px.total); + setSources(srcs); + } catch (e) { + toast.error(`加载失败：${(e as Error).message}`); + } finally { + setLoading(false); + } + }, [page, pageSize]); + + React.useEffect(() => { + void load(); + }, [load]); + + const openCreate = () => { + setEditing(null); + setForm(emptyInput()); + }; + const openEdit = (p: ProxyNode) => { + setEditing(p); + setForm({ + protocol: p.protocol, + host: p.host, + port: p.port, + username: p.username ?? "", + password: "", // 留空 = 不改（后端已打码） + anonymity: p.anonymity ?? "", + region: p.region ?? "", + tags: p.tags, + label: p.label ?? "", + enabled: p.enabled, + }); + }; + + const save = async () => { + if (!form) return; + if (!form.host.trim() || form.port <= 0) { + toast.error("请填写 host 和有效端口"); + return; + } + setSaving(true); + try { + if (editing) { + await api.updateProxy(editing.id, form); + toast.success("已保存"); + } else { + await api.createProxy(form); + toast.success("已添加"); + } + setForm(null); + setEditing(null); + await load(); + } catch (e) { + toast.error(`保存失败：${(e as Error).message}`); + } finally { + setSaving(false); + } + }; + + const remove = async (p: ProxyNode) => { + setProxies((prev) => prev.filter((x) => x.id !== p.id)); + try { + await api.deleteProxy(p.id); + await load(); // 重拉当前页：修正 total 并从后页补位 + } catch (e) { + toast.error(`删除失败：${(e as Error).message}`); + await load(); + } + }; + + const toggleEnabled = async (p: ProxyNode, enabled: boolean) => { + setProxies((prev) => prev.map((x) => (x.id === p.id ? { ...x, enabled } : x))); + try { + await api.updateProxy(p.id, { + protocol: p.protocol, + host: p.host, + port: p.port, + username: p.username, + anonymity: p.anonymity, + region: p.region, + tags: p.tags, + label: p.label, + enabled, + }); + } catch { + await load(); + } + }; + + const check = async (p: ProxyNode) => { + setChecking((prev) => new Set(prev).add(p.id)); + try { + const res = await api.checkProxy(p.id); + toast[res.ok ? "success" : "error"](res.ok ? `可用 · ${res.latency_ms}ms` : `不可用：${res.error}`); + await load(); + } catch (e) { + toast.error(`探活失败：${(e as Error).message}`); + } finally { + setChecking((prev) => { + const next = new Set(prev); + next.delete(p.id); + return next; + }); + } + }; + + const doImport = async () => { + setSaving(true); + try { + const res = await api.importProxies(importText); + toast.success(`已导入 ${res.added} 个${res.invalid.length ? `，${res.invalid.length} 行无法解析` : ""}`); + setImportOpen(false); + setImportText(""); + await load(); + } catch (e) { + toast.error(`导入失败：${(e as Error).message}`); + } finally { + setSaving(false); + } + }; + + const fetchSource = async (name: string) => { + setFetching((prev) => new Set(prev).add(name)); + try { + const res = await api.fetchProxySource(name); + toast.success(`抓取 ${res.fetched} 个，新增 ${res.added} 个`); + await load(); + } catch (e) { + toast.error(`抓取失败：${(e as Error).message}`); + } finally { + setFetching((prev) => { + const next = new Set(prev); + next.delete(name); + return next; + }); + } + }; + + const toggleSource = async (name: string, enabled: boolean) => { + setSources((prev) => prev.map((s) => (s.name === name ? { ...s, enabled } : s))); + try { + await api.setProxySource(name, enabled); + } catch (e) { + toast.error(`切换失败：${(e as Error).message}`); + await load(); + } + }; + + return ( +
+
+

代理池

+

+ 出口代理管理：手动添加/导入可信代理，或开启免费代理源（默认全关）。后台定时验活，选择时优先稳定节点。 +

+
+ + {/* 免费代理源开关 */} + + + 免费代理源 + + 默认全部关闭。免费代理不可信、易被目标封禁——出口默认只走手动/导入的可信代理（见系统配置的安全阀门）。 + + + + {sources.map((s) => ( +
+
+
{s.name}
+
+ {s.last_fetch_at + ? `上次抓取 ${s.last_count} 个${s.last_error ? ` · 错误：${s.last_error}` : ""}` + : "尚未抓取"} +
+
+
+ + toggleSource(s.name, v)} /> +
+
+ ))} + + + + {/* 代理列表 */} + + +
+ 代理节点 + 共 {total} 个 +
+
+ + +
+ + + {loading && ( +
+ 加载中… +
+ )} + {!loading && proxies.length === 0 && ( +
+ 还没有代理，添加或导入，或开启上方免费源。 +
+ )} + {!loading && proxies.length > 0 && ( +
+ + + + 地址 + 协议 + 地区 + 来源 + 健康 + 启用 + 操作 + + + + {proxies.map((p) => ( + + + {p.host}:{p.port} + {p.label ? {p.label} : null} + + + {p.protocol} + + {p.region || "—"} + + {p.trusted ? ( + + 可信 + + ) : ( + + 免费源 + + )} + + + + + + toggleEnabled(p, v)} /> + + +
+ + + +
+ + + ))} + +
+
+ )} + {!loading && total > 0 && ( +
+ { + setPageSize(s); + setPage(1); + }} + /> +
+ )} + + + + {/* 新增/编辑 dialog */} + !o && setForm(null)}> + + + {editing ? "编辑代理" : "添加代理"} + 手动添加的代理标记为可信，可进入主出口轮换。 + + {form && ( +
+
+
+ + setForm({ ...form, protocol: e.target.value as ProxyProtocol })} + > + {PROTOCOLS.map((p) => ( + + {p} + + ))} + +
+
+ + setForm({ ...form, host: e.target.value })} + /> +
+
+
+
+ + setForm({ ...form, port: Number(e.target.value) })} + /> +
+
+ + setForm({ ...form, region: e.target.value })} + /> +
+
+
+
+ + setForm({ ...form, username: e.target.value })} + /> +
+
+ + setForm({ ...form, password: e.target.value })} + /> +
+
+
+ + setForm({ ...form, label: e.target.value })} + /> +
+
+ )} + + + + + + + + {/* 批量导入 dialog */} + + + + 批量导入代理 + + 每行一个：host:port 或 scheme://user:pass@host:port。导入的标记为可信、去重。 + + +