From a6f194aedd521aee7d6557c4fc1612b2c20f8a50 Mon Sep 17 00:00:00 2001 From: krau <71133316+krau@users.noreply.github.com> Date: Tue, 2 Dec 2025 10:11:58 +0800 Subject: [PATCH] feat: add global proxy config --- client/bot/bot.go | 27 ++++-------- client/bot/handlers/link.go | 6 --- client/bot/handlers/media.go | 5 --- client/bot/handlers/parse.go | 5 --- client/bot/handlers/save.go | 6 --- client/bot/handlers/telegraph.go | 6 --- client/user/userclient.go | 26 ++++------- common/utils/netutil/proxy.go | 76 ++++++++++++++++---------------- common/utils/tgutil/net.go | 41 +++++++++++++++++ config/viper.go | 44 ++++++++++++++++++ storage/load.go | 2 +- 11 files changed, 141 insertions(+), 103 deletions(-) create mode 100644 common/utils/tgutil/net.go diff --git a/client/bot/bot.go b/client/bot/bot.go index 85f83d4..a465b21 100644 --- a/client/bot/bot.go +++ b/client/bot/bot.go @@ -9,14 +9,12 @@ import ( "github.com/celestix/gotgproto/ext" "github.com/celestix/gotgproto/sessionMaker" "github.com/charmbracelet/log" - "github.com/gotd/td/telegram/dcs" "github.com/gotd/td/tg" "github.com/krau/SaveAny-Bot/client/bot/handlers" "github.com/krau/SaveAny-Bot/client/middleware" - "github.com/krau/SaveAny-Bot/common/utils/netutil" + "github.com/krau/SaveAny-Bot/common/utils/tgutil" "github.com/krau/SaveAny-Bot/config" "github.com/ncruces/go-sqlite3/gormlite" - "golang.org/x/net/proxy" ) func Init(ctx context.Context) <-chan struct{} { @@ -26,22 +24,15 @@ func Init(ctx context.Context) <-chan struct{} { err error }) shouldRestart := make(chan struct{}) + go func() { - var resolver dcs.Resolver - if config.C().Telegram.Proxy.Enable && config.C().Telegram.Proxy.URL != "" { - dialer, err := netutil.NewProxyDialer(config.C().Telegram.Proxy.URL) - if err != nil { - resultChan <- struct { - client *gotgproto.Client - err error - }{nil, err} - return - } - resolver = dcs.Plain(dcs.PlainOptions{ - Dial: dialer.(proxy.ContextDialer).DialContext, - }) - } else { - resolver = dcs.DefaultResolver() + resolver, err := tgutil.NewConfigProxyResolver() + if err != nil { + resultChan <- struct { + client *gotgproto.Client + err error + }{nil, err} + return } client, err := gotgproto.NewClient( config.C().Telegram.AppID, diff --git a/client/bot/handlers/link.go b/client/bot/handlers/link.go index 107530a..086ce9e 100644 --- a/client/bot/handlers/link.go +++ b/client/bot/handlers/link.go @@ -44,13 +44,7 @@ func handleMessageLink(ctx *ext.Context, update *ext.Update) error { } func handleSilentSaveLink(ctx *ext.Context, update *ext.Update) error { - logger := log.FromContext(ctx) stor := storage.FromContext(ctx) - if stor == nil { - logger.Warn("Context storage is nil") - ctx.Reply(update, ext.ReplyTextString("未找到存储"), nil) - return dispatcher.EndGroups - } replied, files, _, err := shortcut.GetFilesFromUpdateLinkMessageWithReplyEdit(ctx, update) if err != nil { return err diff --git a/client/bot/handlers/media.go b/client/bot/handlers/media.go index bbd9d71..ff9d4c1 100644 --- a/client/bot/handlers/media.go +++ b/client/bot/handlers/media.go @@ -45,11 +45,6 @@ func handleMediaMessage(ctx *ext.Context, update *ext.Update) error { func handleSilentSaveMedia(ctx *ext.Context, update *ext.Update) error { logger := log.FromContext(ctx) stor := storage.FromContext(ctx) - if stor == nil { - logger.Warn("Context storage is nil") - ctx.Reply(update, ext.ReplyTextString("未找到存储"), nil) - return dispatcher.EndGroups - } message := update.EffectiveMessage.Message groupID, isGroup := message.GetGroupedID() if isGroup && groupID != 0 { diff --git a/client/bot/handlers/parse.go b/client/bot/handlers/parse.go index 79b299b..fa4acf4 100644 --- a/client/bot/handlers/parse.go +++ b/client/bot/handlers/parse.go @@ -77,11 +77,6 @@ func handleTextMessage(ctx *ext.Context, u *ext.Update) error { func handleSilentSaveText(ctx *ext.Context, u *ext.Update) error { logger := log.FromContext(ctx) stor := storage.FromContext(ctx) - if stor == nil { - logger.Warn("Context storage is nil") - ctx.Reply(u, ext.ReplyTextString("未找到存储"), nil) - return dispatcher.EndGroups - } text := u.EffectiveMessage.Text if text == "" { return dispatcher.EndGroups diff --git a/client/bot/handlers/save.go b/client/bot/handlers/save.go index eba2f31..4b8db1f 100644 --- a/client/bot/handlers/save.go +++ b/client/bot/handlers/save.go @@ -76,13 +76,7 @@ func handleSilentSaveReplied(ctx *ext.Context, update *ext.Update) error { if len(args) >= 3 { return handleBatchSave(ctx, update, args[1:]) } - logger := log.FromContext(ctx) stor := storage.FromContext(ctx) - if stor == nil { - logger.Warn("Context storage is nil") - ctx.Reply(update, ext.ReplyTextString("未找到存储"), nil) - return dispatcher.EndGroups - } replyTo := update.EffectiveMessage.ReplyToMessage if replyTo == nil || replyTo.Message == nil { ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgSaveHelpText)), nil) diff --git a/client/bot/handlers/telegraph.go b/client/bot/handlers/telegraph.go index ce2227e..421fb6a 100644 --- a/client/bot/handlers/telegraph.go +++ b/client/bot/handlers/telegraph.go @@ -61,13 +61,7 @@ func handleTelegraphUrlMessage(ctx *ext.Context, update *ext.Update) error { } func handleSilentSaveTelegraph(ctx *ext.Context, update *ext.Update) error { - logger := log.FromContext(ctx) stor := storage.FromContext(ctx) - if stor == nil { - logger.Warn("Context storage is nil") - ctx.Reply(update, ext.ReplyTextString("未找到存储"), nil) - return dispatcher.EndGroups - } msg, result, err := shortcut.GetTphPicsFromMessageWithReply(ctx, update) if err != nil { return err diff --git a/client/user/userclient.go b/client/user/userclient.go index cfc5349..a5653ce 100644 --- a/client/user/userclient.go +++ b/client/user/userclient.go @@ -12,14 +12,12 @@ import ( "github.com/celestix/gotgproto/sessionMaker" "github.com/charmbracelet/log" - "github.com/gotd/td/telegram/dcs" "github.com/gotd/td/tg" "github.com/krau/SaveAny-Bot/client/middleware" - "github.com/krau/SaveAny-Bot/common/utils/netutil" + "github.com/krau/SaveAny-Bot/common/utils/tgutil" "github.com/krau/SaveAny-Bot/config" "github.com/krau/SaveAny-Bot/database" "github.com/ncruces/go-sqlite3/gormlite" - "golang.org/x/net/proxy" ) var uc *gotgproto.Client @@ -53,21 +51,13 @@ func Login(ctx context.Context) (*gotgproto.Client, error) { err error }) go func() { - var resolver dcs.Resolver - if config.C().Telegram.Proxy.Enable && config.C().Telegram.Proxy.URL != "" { - dialer, err := netutil.NewProxyDialer(config.C().Telegram.Proxy.URL) - if err != nil { - res <- struct { - client *gotgproto.Client - err error - }{nil, err} - return - } - resolver = dcs.Plain(dcs.PlainOptions{ - Dial: dialer.(proxy.ContextDialer).DialContext, - }) - } else { - resolver = dcs.DefaultResolver() + resolver, err := tgutil.NewConfigProxyResolver() + if err != nil { + res <- struct { + client *gotgproto.Client + err error + }{nil, err} + return } tclient, err := gotgproto.NewClient( config.C().Telegram.AppID, diff --git a/common/utils/netutil/proxy.go b/common/utils/netutil/proxy.go index 473b461..23cd95c 100644 --- a/common/utils/netutil/proxy.go +++ b/common/utils/netutil/proxy.go @@ -7,56 +7,24 @@ import ( "net/http" "net/url" "sync" + "time" "github.com/charmbracelet/log" "github.com/krau/SaveAny-Bot/config" "golang.org/x/net/proxy" ) -func NewProxyDialer(proxyUrl string) (proxy.Dialer, error) { - url, err := url.Parse(proxyUrl) - if err != nil { - return nil, err - } - return proxy.FromURL(url, proxy.Direct) -} - func NewProxyHTTPClient(proxyUrl string) (*http.Client, error) { if proxyUrl == "" { - return &http.Client{ - Transport: &http.Transport{ - Proxy: http.ProxyFromEnvironment, - }, - }, nil + return http.DefaultClient, nil } - - u, err := url.Parse(proxyUrl) + transport, err := NewProxyTransport(proxyUrl) if err != nil { return nil, err } - - switch u.Scheme { - case "http", "https": - return &http.Client{ - Transport: &http.Transport{ - Proxy: http.ProxyURL(u), - }, - }, nil - case "socks5": - dialer, err := proxy.FromURL(u, proxy.Direct) - if err != nil { - return nil, err - } - return &http.Client{ - Transport: &http.Transport{ - DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { - return dialer.Dial(network, addr) - }, - }, - }, nil - default: - return nil, fmt.Errorf("unsupported proxy scheme: %s", u.Scheme) - } + return &http.Client{ + Transport: transport, + }, nil } var ( @@ -76,3 +44,35 @@ func DefaultParserHTTPClient() *http.Client { }) return defaultProxyHttpClient } + +func NewProxyTransport(proxyStr string) (*http.Transport, error) { + proxyURL, err := url.Parse(proxyStr) + if err != nil { + return nil, err + } + transport := &http.Transport{ + ForceAttemptHTTP2: true, + MaxIdleConns: 100, + IdleConnTimeout: 90 * time.Second, + TLSHandshakeTimeout: 10 * time.Second, + ExpectContinueTimeout: 1 * time.Second, + } + switch proxyURL.Scheme { + case "http", "https": + transport.Proxy = http.ProxyURL(proxyURL) + + case "socks5", "socks5h": + dialer, err := proxy.FromURL(proxyURL, proxy.Direct) + if err != nil { + return nil, err + } + transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) { + return dialer.(proxy.ContextDialer).DialContext(ctx, network, addr) + } + + default: + return nil, fmt.Errorf("unsupported proxy type: %s", proxyURL.Scheme) + } + + return transport, nil +} diff --git a/common/utils/tgutil/net.go b/common/utils/tgutil/net.go new file mode 100644 index 0000000..6608cfc --- /dev/null +++ b/common/utils/tgutil/net.go @@ -0,0 +1,41 @@ +package tgutil + +import ( + "net/url" + + "github.com/gotd/td/telegram/dcs" + "github.com/krau/SaveAny-Bot/config" + "golang.org/x/net/proxy" +) + +func newProxyDialer(proxyUrl string) (proxy.Dialer, error) { + url, err := url.Parse(proxyUrl) + if err != nil { + return nil, err + } + return proxy.FromURL(url, proxy.Direct) +} + +func NewConfigProxyResolver() (dcs.Resolver, error) { + resolver := dcs.DefaultResolver() + if config.C().Proxy != "" { + // gloabl proxy, which has lower priority + dialer, err := newProxyDialer(config.C().Proxy) + if err != nil { + return nil, err + } + resolver = dcs.Plain(dcs.PlainOptions{ + Dial: dialer.(proxy.ContextDialer).DialContext, + }) + } + if config.C().Telegram.Proxy.Enable && config.C().Telegram.Proxy.URL != "" { + dialer, err := newProxyDialer(config.C().Telegram.Proxy.URL) + if err != nil { + return nil, err + } + resolver = dcs.Plain(dcs.PlainOptions{ + Dial: dialer.(proxy.ContextDialer).DialContext, + }) + } + return resolver, nil +} diff --git a/config/viper.go b/config/viper.go index 1ed7e8b..f66cebb 100644 --- a/config/viper.go +++ b/config/viper.go @@ -4,13 +4,18 @@ import ( "context" "errors" "fmt" + "net" + "net/http" + "net/url" "strings" + "time" "github.com/duke-git/lancet/v2/slice" "github.com/krau/SaveAny-Bot/common/i18n" "github.com/krau/SaveAny-Bot/common/i18n/i18nk" "github.com/krau/SaveAny-Bot/config/storage" "github.com/spf13/viper" + "golang.org/x/net/proxy" ) type Config struct { @@ -20,6 +25,7 @@ type Config struct { NoCleanCache bool `toml:"no_clean_cache" mapstructure:"no_clean_cache" json:"no_clean_cache"` Threads int `toml:"threads" mapstructure:"threads" json:"threads"` Stream bool `toml:"stream" mapstructure:"stream" json:"stream"` + Proxy string `toml:"proxy" mapstructure:"proxy" json:"proxy"` Cache cacheConfig `toml:"cache" mapstructure:"cache" json:"cache"` Users []userConfig `toml:"users" mapstructure:"users" json:"users"` @@ -147,5 +153,43 @@ func Init(ctx context.Context) error { userStorages[user.ID] = user.Storages } } + if cfg.Proxy != "" { + http.DefaultTransport, err = newProxyTransport(cfg.Proxy) + if err != nil { + return fmt.Errorf("failed to create proxy transport: %w", err) + } + } return nil } + +func newProxyTransport(proxyStr string) (*http.Transport, error) { + proxyURL, err := url.Parse(proxyStr) + if err != nil { + return nil, err + } + transport := &http.Transport{ + ForceAttemptHTTP2: true, + MaxIdleConns: 100, + IdleConnTimeout: 90 * time.Second, + TLSHandshakeTimeout: 10 * time.Second, + ExpectContinueTimeout: 1 * time.Second, + } + switch proxyURL.Scheme { + case "http", "https": + transport.Proxy = http.ProxyURL(proxyURL) + + case "socks5", "socks5h": + dialer, err := proxy.FromURL(proxyURL, proxy.Direct) + if err != nil { + return nil, err + } + transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) { + return dialer.(proxy.ContextDialer).DialContext(ctx, network, addr) + } + + default: + return nil, fmt.Errorf("unsupported proxy type: %s", proxyURL.Scheme) + } + + return transport, nil +} diff --git a/storage/load.go b/storage/load.go index 4d67f01..ed9c845 100644 --- a/storage/load.go +++ b/storage/load.go @@ -40,7 +40,7 @@ func GetStorageByUserIDAndName(ctx context.Context, chatID int64, name string) ( } if !config.C().HasStorage(chatID, name) { - return nil, fmt.Errorf("没有找到用户 %d 的存储 %s", chatID, name) + return nil, fmt.Errorf("no storage %s for user %d", name, chatID) } return getStorageByName(ctx, name)