merge: 合并 upstream/main 并保留本地图片计费功能
This commit is contained in:
@@ -67,6 +67,7 @@ func (r *accountRepository) Create(ctx context.Context, account *service.Account
|
||||
|
||||
builder := r.client.Account.Create().
|
||||
SetName(account.Name).
|
||||
SetNillableNotes(account.Notes).
|
||||
SetPlatform(account.Platform).
|
||||
SetType(account.Type).
|
||||
SetCredentials(normalizeJSONMap(account.Credentials)).
|
||||
@@ -270,6 +271,7 @@ func (r *accountRepository) Update(ctx context.Context, account *service.Account
|
||||
|
||||
builder := r.client.Account.UpdateOneID(account.ID).
|
||||
SetName(account.Name).
|
||||
SetNillableNotes(account.Notes).
|
||||
SetPlatform(account.Platform).
|
||||
SetType(account.Type).
|
||||
SetCredentials(normalizeJSONMap(account.Credentials)).
|
||||
@@ -320,6 +322,9 @@ func (r *accountRepository) Update(ctx context.Context, account *service.Account
|
||||
} else {
|
||||
builder.ClearSessionWindowStatus()
|
||||
}
|
||||
if account.Notes == nil {
|
||||
builder.ClearNotes()
|
||||
}
|
||||
|
||||
updated, err := builder.Save(ctx)
|
||||
if err != nil {
|
||||
@@ -768,9 +773,14 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
idx++
|
||||
}
|
||||
if updates.ProxyID != nil {
|
||||
setClauses = append(setClauses, "proxy_id = $"+itoa(idx))
|
||||
args = append(args, *updates.ProxyID)
|
||||
idx++
|
||||
// 0 表示清除代理(前端发送 0 而不是 null 来表达清除意图)
|
||||
if *updates.ProxyID == 0 {
|
||||
setClauses = append(setClauses, "proxy_id = NULL")
|
||||
} else {
|
||||
setClauses = append(setClauses, "proxy_id = $"+itoa(idx))
|
||||
args = append(args, *updates.ProxyID)
|
||||
idx++
|
||||
}
|
||||
}
|
||||
if updates.Concurrency != nil {
|
||||
setClauses = append(setClauses, "concurrency = $"+itoa(idx))
|
||||
@@ -1065,6 +1075,7 @@ func accountEntityToService(m *dbent.Account) *service.Account {
|
||||
return &service.Account{
|
||||
ID: m.ID,
|
||||
Name: m.Name,
|
||||
Notes: m.Notes,
|
||||
Platform: m.Platform,
|
||||
Type: m.Type,
|
||||
Credentials: copyJSONMap(m.Credentials),
|
||||
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/oauth"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/logredact"
|
||||
|
||||
"github.com/imroc/req/v3"
|
||||
)
|
||||
@@ -54,7 +55,7 @@ func (s *claudeOAuthService) GetOrganizationUUID(ctx context.Context, sessionKey
|
||||
return "", fmt.Errorf("request failed: %w", err)
|
||||
}
|
||||
|
||||
log.Printf("[OAuth] Step 1 Response - Status: %d, Body: %s", resp.StatusCode, resp.String())
|
||||
log.Printf("[OAuth] Step 1 Response - Status: %d", resp.StatusCode)
|
||||
|
||||
if !resp.IsSuccessState() {
|
||||
return "", fmt.Errorf("failed to get organizations: status %d, body: %s", resp.StatusCode, resp.String())
|
||||
@@ -84,8 +85,8 @@ func (s *claudeOAuthService) GetAuthorizationCode(ctx context.Context, sessionKe
|
||||
"code_challenge_method": "S256",
|
||||
}
|
||||
|
||||
reqBodyJSON, _ := json.Marshal(reqBody)
|
||||
log.Printf("[OAuth] Step 2: Getting authorization code from %s", authURL)
|
||||
reqBodyJSON, _ := json.Marshal(logredact.RedactMap(reqBody))
|
||||
log.Printf("[OAuth] Step 2 Request Body: %s", string(reqBodyJSON))
|
||||
|
||||
var result struct {
|
||||
@@ -113,7 +114,7 @@ func (s *claudeOAuthService) GetAuthorizationCode(ctx context.Context, sessionKe
|
||||
return "", fmt.Errorf("request failed: %w", err)
|
||||
}
|
||||
|
||||
log.Printf("[OAuth] Step 2 Response - Status: %d, Body: %s", resp.StatusCode, resp.String())
|
||||
log.Printf("[OAuth] Step 2 Response - Status: %d, Body: %s", resp.StatusCode, logredact.RedactJSON(resp.Bytes()))
|
||||
|
||||
if !resp.IsSuccessState() {
|
||||
return "", fmt.Errorf("failed to get authorization code: status %d, body: %s", resp.StatusCode, resp.String())
|
||||
@@ -141,7 +142,7 @@ func (s *claudeOAuthService) GetAuthorizationCode(ctx context.Context, sessionKe
|
||||
fullCode = authCode + "#" + responseState
|
||||
}
|
||||
|
||||
log.Printf("[OAuth] Step 2 SUCCESS - Got authorization code: %s...", prefix(authCode, 20))
|
||||
log.Printf("[OAuth] Step 2 SUCCESS - Got authorization code")
|
||||
return fullCode, nil
|
||||
}
|
||||
|
||||
@@ -173,8 +174,8 @@ func (s *claudeOAuthService) ExchangeCodeForToken(ctx context.Context, code, cod
|
||||
reqBody["expires_in"] = 31536000 // 365 * 24 * 60 * 60 seconds
|
||||
}
|
||||
|
||||
reqBodyJSON, _ := json.Marshal(reqBody)
|
||||
log.Printf("[OAuth] Step 3: Exchanging code for token at %s", s.tokenURL)
|
||||
reqBodyJSON, _ := json.Marshal(logredact.RedactMap(reqBody))
|
||||
log.Printf("[OAuth] Step 3 Request Body: %s", string(reqBodyJSON))
|
||||
|
||||
var tokenResp oauth.TokenResponse
|
||||
@@ -191,7 +192,7 @@ func (s *claudeOAuthService) ExchangeCodeForToken(ctx context.Context, code, cod
|
||||
return nil, fmt.Errorf("request failed: %w", err)
|
||||
}
|
||||
|
||||
log.Printf("[OAuth] Step 3 Response - Status: %d, Body: %s", resp.StatusCode, resp.String())
|
||||
log.Printf("[OAuth] Step 3 Response - Status: %d, Body: %s", resp.StatusCode, logredact.RedactJSON(resp.Bytes()))
|
||||
|
||||
if !resp.IsSuccessState() {
|
||||
return nil, fmt.Errorf("token exchange failed: status %d, body: %s", resp.StatusCode, resp.String())
|
||||
@@ -245,13 +246,3 @@ func createReqClient(proxyURL string) *req.Client {
|
||||
|
||||
return client
|
||||
}
|
||||
|
||||
func prefix(s string, n int) string {
|
||||
if n <= 0 {
|
||||
return ""
|
||||
}
|
||||
if len(s) <= n {
|
||||
return s
|
||||
}
|
||||
return s[:n]
|
||||
}
|
||||
|
||||
@@ -15,7 +15,8 @@ import (
|
||||
const defaultClaudeUsageURL = "https://api.anthropic.com/api/oauth/usage"
|
||||
|
||||
type claudeUsageService struct {
|
||||
usageURL string
|
||||
usageURL string
|
||||
allowPrivateHosts bool
|
||||
}
|
||||
|
||||
func NewClaudeUsageFetcher() service.ClaudeUsageFetcher {
|
||||
@@ -24,8 +25,10 @@ func NewClaudeUsageFetcher() service.ClaudeUsageFetcher {
|
||||
|
||||
func (s *claudeUsageService) FetchUsage(ctx context.Context, accessToken, proxyURL string) (*service.ClaudeUsageResponse, error) {
|
||||
client, err := httpclient.GetClient(httpclient.Options{
|
||||
ProxyURL: proxyURL,
|
||||
Timeout: 30 * time.Second,
|
||||
ProxyURL: proxyURL,
|
||||
Timeout: 30 * time.Second,
|
||||
ValidateResolvedIP: true,
|
||||
AllowPrivateHosts: s.allowPrivateHosts,
|
||||
})
|
||||
if err != nil {
|
||||
client = &http.Client{Timeout: 30 * time.Second}
|
||||
|
||||
@@ -45,7 +45,10 @@ func (s *ClaudeUsageServiceSuite) TestFetchUsage_Success() {
|
||||
}`)
|
||||
}))
|
||||
|
||||
s.fetcher = &claudeUsageService{usageURL: s.srv.URL}
|
||||
s.fetcher = &claudeUsageService{
|
||||
usageURL: s.srv.URL,
|
||||
allowPrivateHosts: true,
|
||||
}
|
||||
|
||||
resp, err := s.fetcher.FetchUsage(context.Background(), "at", "://bad-proxy-url")
|
||||
require.NoError(s.T(), err, "FetchUsage")
|
||||
@@ -64,7 +67,10 @@ func (s *ClaudeUsageServiceSuite) TestFetchUsage_NonOK() {
|
||||
_, _ = io.WriteString(w, "nope")
|
||||
}))
|
||||
|
||||
s.fetcher = &claudeUsageService{usageURL: s.srv.URL}
|
||||
s.fetcher = &claudeUsageService{
|
||||
usageURL: s.srv.URL,
|
||||
allowPrivateHosts: true,
|
||||
}
|
||||
|
||||
_, err := s.fetcher.FetchUsage(context.Background(), "at", "")
|
||||
require.Error(s.T(), err)
|
||||
@@ -78,7 +84,10 @@ func (s *ClaudeUsageServiceSuite) TestFetchUsage_BadJSON() {
|
||||
_, _ = io.WriteString(w, "not-json")
|
||||
}))
|
||||
|
||||
s.fetcher = &claudeUsageService{usageURL: s.srv.URL}
|
||||
s.fetcher = &claudeUsageService{
|
||||
usageURL: s.srv.URL,
|
||||
allowPrivateHosts: true,
|
||||
}
|
||||
|
||||
_, err := s.fetcher.FetchUsage(context.Background(), "at", "")
|
||||
require.Error(s.T(), err)
|
||||
@@ -91,7 +100,10 @@ func (s *ClaudeUsageServiceSuite) TestFetchUsage_ContextCancel() {
|
||||
<-r.Context().Done()
|
||||
}))
|
||||
|
||||
s.fetcher = &claudeUsageService{usageURL: s.srv.URL}
|
||||
s.fetcher = &claudeUsageService{
|
||||
usageURL: s.srv.URL,
|
||||
allowPrivateHosts: true,
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel() // Cancel immediately
|
||||
|
||||
@@ -56,7 +56,7 @@ func InitEnt(cfg *config.Config) (*ent.Client, *sql.DB, error) {
|
||||
// 确保数据库 schema 已准备就绪。
|
||||
// SQL 迁移文件是 schema 的权威来源(source of truth)。
|
||||
// 这种方式比 Ent 的自动迁移更可控,支持复杂的迁移场景。
|
||||
migrationCtx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
migrationCtx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
|
||||
defer cancel()
|
||||
if err := applyMigrationsFS(migrationCtx, drv.DB(), migrations.FS); err != nil {
|
||||
_ = drv.Close() // 迁移失败时关闭驱动,避免资源泄露
|
||||
|
||||
@@ -14,18 +14,23 @@ import (
|
||||
)
|
||||
|
||||
type githubReleaseClient struct {
|
||||
httpClient *http.Client
|
||||
httpClient *http.Client
|
||||
allowPrivateHosts bool
|
||||
}
|
||||
|
||||
func NewGitHubReleaseClient() service.GitHubReleaseClient {
|
||||
allowPrivate := false
|
||||
sharedClient, err := httpclient.GetClient(httpclient.Options{
|
||||
Timeout: 30 * time.Second,
|
||||
Timeout: 30 * time.Second,
|
||||
ValidateResolvedIP: true,
|
||||
AllowPrivateHosts: allowPrivate,
|
||||
})
|
||||
if err != nil {
|
||||
sharedClient = &http.Client{Timeout: 30 * time.Second}
|
||||
}
|
||||
return &githubReleaseClient{
|
||||
httpClient: sharedClient,
|
||||
httpClient: sharedClient,
|
||||
allowPrivateHosts: allowPrivate,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -64,7 +69,9 @@ func (c *githubReleaseClient) DownloadFile(ctx context.Context, url, dest string
|
||||
}
|
||||
|
||||
downloadClient, err := httpclient.GetClient(httpclient.Options{
|
||||
Timeout: 10 * time.Minute,
|
||||
Timeout: 10 * time.Minute,
|
||||
ValidateResolvedIP: true,
|
||||
AllowPrivateHosts: c.allowPrivateHosts,
|
||||
})
|
||||
if err != nil {
|
||||
downloadClient = &http.Client{Timeout: 10 * time.Minute}
|
||||
|
||||
@@ -37,6 +37,13 @@ func (t *testTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return http.DefaultTransport.RoundTrip(newReq)
|
||||
}
|
||||
|
||||
func newTestGitHubReleaseClient() *githubReleaseClient {
|
||||
return &githubReleaseClient{
|
||||
httpClient: &http.Client{},
|
||||
allowPrivateHosts: true,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *GitHubReleaseServiceSuite) SetupTest() {
|
||||
s.tempDir = s.T().TempDir()
|
||||
}
|
||||
@@ -55,9 +62,7 @@ func (s *GitHubReleaseServiceSuite) TestDownloadFile_EnforcesMaxSize_ContentLeng
|
||||
_, _ = w.Write(bytes.Repeat([]byte("a"), 100))
|
||||
}))
|
||||
|
||||
client, ok := NewGitHubReleaseClient().(*githubReleaseClient)
|
||||
require.True(s.T(), ok, "type assertion failed")
|
||||
s.client = client
|
||||
s.client = newTestGitHubReleaseClient()
|
||||
|
||||
dest := filepath.Join(s.tempDir, "file1.bin")
|
||||
err := s.client.DownloadFile(context.Background(), s.srv.URL, dest, 10)
|
||||
@@ -82,9 +87,7 @@ func (s *GitHubReleaseServiceSuite) TestDownloadFile_EnforcesMaxSize_Chunked() {
|
||||
}
|
||||
}))
|
||||
|
||||
client, ok := NewGitHubReleaseClient().(*githubReleaseClient)
|
||||
require.True(s.T(), ok, "type assertion failed")
|
||||
s.client = client
|
||||
s.client = newTestGitHubReleaseClient()
|
||||
|
||||
dest := filepath.Join(s.tempDir, "file2.bin")
|
||||
err := s.client.DownloadFile(context.Background(), s.srv.URL, dest, 10)
|
||||
@@ -108,9 +111,7 @@ func (s *GitHubReleaseServiceSuite) TestDownloadFile_Success() {
|
||||
}
|
||||
}))
|
||||
|
||||
client, ok := NewGitHubReleaseClient().(*githubReleaseClient)
|
||||
require.True(s.T(), ok, "type assertion failed")
|
||||
s.client = client
|
||||
s.client = newTestGitHubReleaseClient()
|
||||
|
||||
dest := filepath.Join(s.tempDir, "file3.bin")
|
||||
err := s.client.DownloadFile(context.Background(), s.srv.URL, dest, 200)
|
||||
@@ -127,9 +128,7 @@ func (s *GitHubReleaseServiceSuite) TestDownloadFile_404() {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}))
|
||||
|
||||
client, ok := NewGitHubReleaseClient().(*githubReleaseClient)
|
||||
require.True(s.T(), ok, "type assertion failed")
|
||||
s.client = client
|
||||
s.client = newTestGitHubReleaseClient()
|
||||
|
||||
dest := filepath.Join(s.tempDir, "notfound.bin")
|
||||
err := s.client.DownloadFile(context.Background(), s.srv.URL, dest, 100)
|
||||
@@ -145,9 +144,7 @@ func (s *GitHubReleaseServiceSuite) TestFetchChecksumFile_Success() {
|
||||
_, _ = w.Write([]byte("sum"))
|
||||
}))
|
||||
|
||||
client, ok := NewGitHubReleaseClient().(*githubReleaseClient)
|
||||
require.True(s.T(), ok, "type assertion failed")
|
||||
s.client = client
|
||||
s.client = newTestGitHubReleaseClient()
|
||||
|
||||
body, err := s.client.FetchChecksumFile(context.Background(), s.srv.URL)
|
||||
require.NoError(s.T(), err, "FetchChecksumFile")
|
||||
@@ -159,9 +156,7 @@ func (s *GitHubReleaseServiceSuite) TestFetchChecksumFile_Non200() {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
}))
|
||||
|
||||
client, ok := NewGitHubReleaseClient().(*githubReleaseClient)
|
||||
require.True(s.T(), ok, "type assertion failed")
|
||||
s.client = client
|
||||
s.client = newTestGitHubReleaseClient()
|
||||
|
||||
_, err := s.client.FetchChecksumFile(context.Background(), s.srv.URL)
|
||||
require.Error(s.T(), err, "expected error for non-200")
|
||||
@@ -172,9 +167,7 @@ func (s *GitHubReleaseServiceSuite) TestDownloadFile_ContextCancel() {
|
||||
<-r.Context().Done()
|
||||
}))
|
||||
|
||||
client, ok := NewGitHubReleaseClient().(*githubReleaseClient)
|
||||
require.True(s.T(), ok, "type assertion failed")
|
||||
s.client = client
|
||||
s.client = newTestGitHubReleaseClient()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
@@ -185,9 +178,7 @@ func (s *GitHubReleaseServiceSuite) TestDownloadFile_ContextCancel() {
|
||||
}
|
||||
|
||||
func (s *GitHubReleaseServiceSuite) TestDownloadFile_InvalidURL() {
|
||||
client, ok := NewGitHubReleaseClient().(*githubReleaseClient)
|
||||
require.True(s.T(), ok, "type assertion failed")
|
||||
s.client = client
|
||||
s.client = newTestGitHubReleaseClient()
|
||||
|
||||
dest := filepath.Join(s.tempDir, "invalid.bin")
|
||||
err := s.client.DownloadFile(context.Background(), "://invalid-url", dest, 100)
|
||||
@@ -200,9 +191,7 @@ func (s *GitHubReleaseServiceSuite) TestDownloadFile_InvalidDestPath() {
|
||||
_, _ = w.Write([]byte("content"))
|
||||
}))
|
||||
|
||||
client, ok := NewGitHubReleaseClient().(*githubReleaseClient)
|
||||
require.True(s.T(), ok, "type assertion failed")
|
||||
s.client = client
|
||||
s.client = newTestGitHubReleaseClient()
|
||||
|
||||
// Use a path that cannot be created (directory doesn't exist)
|
||||
dest := filepath.Join(s.tempDir, "nonexistent", "subdir", "file.bin")
|
||||
@@ -211,9 +200,7 @@ func (s *GitHubReleaseServiceSuite) TestDownloadFile_InvalidDestPath() {
|
||||
}
|
||||
|
||||
func (s *GitHubReleaseServiceSuite) TestFetchChecksumFile_InvalidURL() {
|
||||
client, ok := NewGitHubReleaseClient().(*githubReleaseClient)
|
||||
require.True(s.T(), ok, "type assertion failed")
|
||||
s.client = client
|
||||
s.client = newTestGitHubReleaseClient()
|
||||
|
||||
_, err := s.client.FetchChecksumFile(context.Background(), "://invalid-url")
|
||||
require.Error(s.T(), err, "expected error for invalid URL")
|
||||
@@ -247,6 +234,7 @@ func (s *GitHubReleaseServiceSuite) TestFetchLatestRelease_Success() {
|
||||
httpClient: &http.Client{
|
||||
Transport: &testTransport{testServerURL: s.srv.URL},
|
||||
},
|
||||
allowPrivateHosts: true,
|
||||
}
|
||||
|
||||
release, err := s.client.FetchLatestRelease(context.Background(), "test/repo")
|
||||
@@ -266,6 +254,7 @@ func (s *GitHubReleaseServiceSuite) TestFetchLatestRelease_Non200() {
|
||||
httpClient: &http.Client{
|
||||
Transport: &testTransport{testServerURL: s.srv.URL},
|
||||
},
|
||||
allowPrivateHosts: true,
|
||||
}
|
||||
|
||||
_, err := s.client.FetchLatestRelease(context.Background(), "test/repo")
|
||||
@@ -283,6 +272,7 @@ func (s *GitHubReleaseServiceSuite) TestFetchLatestRelease_InvalidJSON() {
|
||||
httpClient: &http.Client{
|
||||
Transport: &testTransport{testServerURL: s.srv.URL},
|
||||
},
|
||||
allowPrivateHosts: true,
|
||||
}
|
||||
|
||||
_, err := s.client.FetchLatestRelease(context.Background(), "test/repo")
|
||||
@@ -298,6 +288,7 @@ func (s *GitHubReleaseServiceSuite) TestFetchLatestRelease_ContextCancel() {
|
||||
httpClient: &http.Client{
|
||||
Transport: &testTransport{testServerURL: s.srv.URL},
|
||||
},
|
||||
allowPrivateHosts: true,
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
@@ -312,9 +303,7 @@ func (s *GitHubReleaseServiceSuite) TestFetchChecksumFile_ContextCancel() {
|
||||
<-r.Context().Done()
|
||||
}))
|
||||
|
||||
client, ok := NewGitHubReleaseClient().(*githubReleaseClient)
|
||||
require.True(s.T(), ok, "type assertion failed")
|
||||
s.client = client
|
||||
s.client = newTestGitHubReleaseClient()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
|
||||
)
|
||||
|
||||
// 默认配置常量
|
||||
@@ -30,9 +31,9 @@ const (
|
||||
// defaultMaxConnsPerHost: 默认每主机最大连接数(含活跃连接)
|
||||
// 达到上限后新请求会等待,而非无限创建连接
|
||||
defaultMaxConnsPerHost = 240
|
||||
// defaultIdleConnTimeout: 默认空闲连接超时时间(5分钟)
|
||||
// 超时后连接会被关闭,释放系统资源
|
||||
defaultIdleConnTimeout = 300 * time.Second
|
||||
// defaultIdleConnTimeout: 默认空闲连接超时时间(90秒)
|
||||
// 超时后连接会被关闭,释放系统资源(建议小于上游 LB 超时)
|
||||
defaultIdleConnTimeout = 90 * time.Second
|
||||
// defaultResponseHeaderTimeout: 默认等待响应头超时时间(5分钟)
|
||||
// LLM 请求可能排队较久,需要较长超时
|
||||
defaultResponseHeaderTimeout = 300 * time.Second
|
||||
@@ -120,6 +121,10 @@ func NewHTTPUpstream(cfg *config.Config) service.HTTPUpstream {
|
||||
// - 调用方必须关闭 resp.Body,否则会导致 inFlight 计数泄漏
|
||||
// - inFlight > 0 的客户端不会被淘汰,确保活跃请求不被中断
|
||||
func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
|
||||
if err := s.validateRequestHost(req); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 获取或创建对应的客户端,并标记请求占用
|
||||
entry, err := s.acquireClient(proxyURL, accountID, accountConcurrency)
|
||||
if err != nil {
|
||||
@@ -145,6 +150,40 @@ func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID i
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (s *httpUpstreamService) shouldValidateResolvedIP() bool {
|
||||
if s.cfg == nil {
|
||||
return false
|
||||
}
|
||||
if !s.cfg.Security.URLAllowlist.Enabled {
|
||||
return false
|
||||
}
|
||||
return !s.cfg.Security.URLAllowlist.AllowPrivateHosts
|
||||
}
|
||||
|
||||
func (s *httpUpstreamService) validateRequestHost(req *http.Request) error {
|
||||
if !s.shouldValidateResolvedIP() {
|
||||
return nil
|
||||
}
|
||||
if req == nil || req.URL == nil {
|
||||
return errors.New("request url is nil")
|
||||
}
|
||||
host := strings.TrimSpace(req.URL.Hostname())
|
||||
if host == "" {
|
||||
return errors.New("request host is empty")
|
||||
}
|
||||
if err := urlvalidator.ValidateResolvedIP(host); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *httpUpstreamService) redirectChecker(req *http.Request, via []*http.Request) error {
|
||||
if len(via) >= 10 {
|
||||
return errors.New("stopped after 10 redirects")
|
||||
}
|
||||
return s.validateRequestHost(req)
|
||||
}
|
||||
|
||||
// acquireClient 获取或创建客户端,并标记为进行中请求
|
||||
// 用于请求路径,避免在获取后被淘汰
|
||||
func (s *httpUpstreamService) acquireClient(proxyURL string, accountID int64, accountConcurrency int) (*upstreamClientEntry, error) {
|
||||
@@ -232,6 +271,9 @@ func (s *httpUpstreamService) getClientEntry(proxyURL string, accountID int64, a
|
||||
return nil, fmt.Errorf("build transport: %w", err)
|
||||
}
|
||||
client := &http.Client{Transport: transport}
|
||||
if s.shouldValidateResolvedIP() {
|
||||
client.CheckRedirect = s.redirectChecker
|
||||
}
|
||||
entry := &upstreamClientEntry{
|
||||
client: client,
|
||||
proxyKey: proxyKey,
|
||||
|
||||
@@ -22,7 +22,13 @@ type HTTPUpstreamSuite struct {
|
||||
// SetupTest 每个测试用例执行前的初始化
|
||||
// 创建空配置,各测试用例可按需覆盖
|
||||
func (s *HTTPUpstreamSuite) SetupTest() {
|
||||
s.cfg = &config.Config{}
|
||||
s.cfg = &config.Config{
|
||||
Security: config.SecurityConfig{
|
||||
URLAllowlist: config.URLAllowlistConfig{
|
||||
AllowPrivateHosts: true,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// newService 创建测试用的 httpUpstreamService 实例
|
||||
|
||||
@@ -26,6 +26,7 @@ func TestMigrationsRunner_IsIdempotent_AndSchemaIsUpToDate(t *testing.T) {
|
||||
requireColumn(t, tx, "users", "notes", "text", 0, false)
|
||||
|
||||
// accounts: schedulable and rate-limit fields
|
||||
requireColumn(t, tx, "accounts", "notes", "text", 0, true)
|
||||
requireColumn(t, tx, "accounts", "schedulable", "boolean", 0, false)
|
||||
requireColumn(t, tx, "accounts", "rate_limited_at", "timestamp with time zone", 0, true)
|
||||
requireColumn(t, tx, "accounts", "rate_limit_reset_at", "timestamp with time zone", 0, true)
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
@@ -16,9 +17,17 @@ type pricingRemoteClient struct {
|
||||
httpClient *http.Client
|
||||
}
|
||||
|
||||
func NewPricingRemoteClient() service.PricingRemoteClient {
|
||||
func NewPricingRemoteClient(cfg *config.Config) service.PricingRemoteClient {
|
||||
allowPrivate := false
|
||||
validateResolvedIP := true
|
||||
if cfg != nil {
|
||||
allowPrivate = cfg.Security.URLAllowlist.AllowPrivateHosts
|
||||
validateResolvedIP = cfg.Security.URLAllowlist.Enabled
|
||||
}
|
||||
sharedClient, err := httpclient.GetClient(httpclient.Options{
|
||||
Timeout: 30 * time.Second,
|
||||
Timeout: 30 * time.Second,
|
||||
ValidateResolvedIP: validateResolvedIP,
|
||||
AllowPrivateHosts: allowPrivate,
|
||||
})
|
||||
if err != nil {
|
||||
sharedClient = &http.Client{Timeout: 30 * time.Second}
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
@@ -19,7 +20,13 @@ type PricingServiceSuite struct {
|
||||
|
||||
func (s *PricingServiceSuite) SetupTest() {
|
||||
s.ctx = context.Background()
|
||||
client, ok := NewPricingRemoteClient().(*pricingRemoteClient)
|
||||
client, ok := NewPricingRemoteClient(&config.Config{
|
||||
Security: config.SecurityConfig{
|
||||
URLAllowlist: config.URLAllowlistConfig{
|
||||
AllowPrivateHosts: true,
|
||||
},
|
||||
},
|
||||
}).(*pricingRemoteClient)
|
||||
require.True(s.T(), ok, "type assertion failed")
|
||||
s.client = client
|
||||
}
|
||||
|
||||
@@ -5,28 +5,52 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
func NewProxyExitInfoProber() service.ProxyExitInfoProber {
|
||||
return &proxyProbeService{ipInfoURL: defaultIPInfoURL}
|
||||
func NewProxyExitInfoProber(cfg *config.Config) service.ProxyExitInfoProber {
|
||||
insecure := false
|
||||
allowPrivate := false
|
||||
validateResolvedIP := true
|
||||
if cfg != nil {
|
||||
insecure = cfg.Security.ProxyProbe.InsecureSkipVerify
|
||||
allowPrivate = cfg.Security.URLAllowlist.AllowPrivateHosts
|
||||
validateResolvedIP = cfg.Security.URLAllowlist.Enabled
|
||||
}
|
||||
if insecure {
|
||||
log.Printf("[ProxyProbe] Warning: TLS verification is disabled for proxy probing.")
|
||||
}
|
||||
return &proxyProbeService{
|
||||
ipInfoURL: defaultIPInfoURL,
|
||||
insecureSkipVerify: insecure,
|
||||
allowPrivateHosts: allowPrivate,
|
||||
validateResolvedIP: validateResolvedIP,
|
||||
}
|
||||
}
|
||||
|
||||
const defaultIPInfoURL = "https://ipinfo.io/json"
|
||||
|
||||
type proxyProbeService struct {
|
||||
ipInfoURL string
|
||||
ipInfoURL string
|
||||
insecureSkipVerify bool
|
||||
allowPrivateHosts bool
|
||||
validateResolvedIP bool
|
||||
}
|
||||
|
||||
func (s *proxyProbeService) ProbeProxy(ctx context.Context, proxyURL string) (*service.ProxyExitInfo, int64, error) {
|
||||
client, err := httpclient.GetClient(httpclient.Options{
|
||||
ProxyURL: proxyURL,
|
||||
Timeout: 15 * time.Second,
|
||||
InsecureSkipVerify: true,
|
||||
InsecureSkipVerify: s.insecureSkipVerify,
|
||||
ProxyStrict: true,
|
||||
ValidateResolvedIP: s.validateResolvedIP,
|
||||
AllowPrivateHosts: s.allowPrivateHosts,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("failed to create proxy client: %w", err)
|
||||
|
||||
@@ -20,7 +20,10 @@ type ProxyProbeServiceSuite struct {
|
||||
|
||||
func (s *ProxyProbeServiceSuite) SetupTest() {
|
||||
s.ctx = context.Background()
|
||||
s.prober = &proxyProbeService{ipInfoURL: "http://ipinfo.test/json"}
|
||||
s.prober = &proxyProbeService{
|
||||
ipInfoURL: "http://ipinfo.test/json",
|
||||
allowPrivateHosts: true,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ProxyProbeServiceSuite) TearDownTest() {
|
||||
|
||||
@@ -22,7 +22,8 @@ type turnstileVerifier struct {
|
||||
|
||||
func NewTurnstileVerifier() service.TurnstileVerifier {
|
||||
sharedClient, err := httpclient.GetClient(httpclient.Options{
|
||||
Timeout: 10 * time.Second,
|
||||
Timeout: 10 * time.Second,
|
||||
ValidateResolvedIP: true,
|
||||
})
|
||||
if err != nil {
|
||||
sharedClient = &http.Client{Timeout: 10 * time.Second}
|
||||
|
||||
@@ -329,17 +329,20 @@ func (r *userRepository) UpdateBalance(ctx context.Context, id int64, amount flo
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeductBalance 扣除用户余额
|
||||
// 透支策略:允许余额变为负数,确保当前请求能够完成
|
||||
// 中间件会阻止余额 <= 0 的用户发起后续请求
|
||||
func (r *userRepository) DeductBalance(ctx context.Context, id int64, amount float64) error {
|
||||
client := clientFromContext(ctx, r.client)
|
||||
n, err := client.User.Update().
|
||||
Where(dbuser.IDEQ(id), dbuser.BalanceGTE(amount)).
|
||||
Where(dbuser.IDEQ(id)).
|
||||
AddBalance(-amount).
|
||||
Save(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n == 0 {
|
||||
return service.ErrInsufficientBalance
|
||||
return service.ErrUserNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -290,9 +290,14 @@ func (s *UserRepoSuite) TestDeductBalance() {
|
||||
func (s *UserRepoSuite) TestDeductBalance_InsufficientFunds() {
|
||||
user := s.mustCreateUser(&service.User{Email: "insuf@test.com", Balance: 5})
|
||||
|
||||
// 透支策略:允许扣除超过余额的金额
|
||||
err := s.repo.DeductBalance(s.ctx, user.ID, 999)
|
||||
s.Require().Error(err, "expected error for insufficient balance")
|
||||
s.Require().ErrorIs(err, service.ErrInsufficientBalance)
|
||||
s.Require().NoError(err, "DeductBalance should allow overdraft")
|
||||
|
||||
// 验证余额变为负数
|
||||
got, err := s.repo.GetByID(s.ctx, user.ID)
|
||||
s.Require().NoError(err)
|
||||
s.Require().InDelta(-994.0, got.Balance, 1e-6, "Balance should be negative after overdraft")
|
||||
}
|
||||
|
||||
func (s *UserRepoSuite) TestDeductBalance_ExactAmount() {
|
||||
@@ -306,6 +311,19 @@ func (s *UserRepoSuite) TestDeductBalance_ExactAmount() {
|
||||
s.Require().InDelta(0.0, got.Balance, 1e-6)
|
||||
}
|
||||
|
||||
func (s *UserRepoSuite) TestDeductBalance_AllowsOverdraft() {
|
||||
user := s.mustCreateUser(&service.User{Email: "overdraft@test.com", Balance: 5.0})
|
||||
|
||||
// 扣除超过余额的金额 - 应该成功
|
||||
err := s.repo.DeductBalance(s.ctx, user.ID, 10.0)
|
||||
s.Require().NoError(err, "DeductBalance should allow overdraft")
|
||||
|
||||
// 验证余额为负
|
||||
got, err := s.repo.GetByID(s.ctx, user.ID)
|
||||
s.Require().NoError(err)
|
||||
s.Require().InDelta(-5.0, got.Balance, 1e-6, "Balance should be -5.0 after overdraft")
|
||||
}
|
||||
|
||||
// --- Concurrency ---
|
||||
|
||||
func (s *UserRepoSuite) TestUpdateConcurrency() {
|
||||
@@ -477,9 +495,12 @@ func (s *UserRepoSuite) TestCRUD_And_Filters_And_AtomicUpdates() {
|
||||
s.Require().NoError(err, "GetByID after DeductBalance")
|
||||
s.Require().InDelta(7.5, got4.Balance, 1e-6)
|
||||
|
||||
// 透支策略:允许扣除超过余额的金额
|
||||
err = s.repo.DeductBalance(s.ctx, user1.ID, 999)
|
||||
s.Require().Error(err, "DeductBalance expected error for insufficient balance")
|
||||
s.Require().ErrorIs(err, service.ErrInsufficientBalance, "DeductBalance unexpected error")
|
||||
s.Require().NoError(err, "DeductBalance should allow overdraft")
|
||||
gotOverdraft, err := s.repo.GetByID(s.ctx, user1.ID)
|
||||
s.Require().NoError(err, "GetByID after overdraft")
|
||||
s.Require().Less(gotOverdraft.Balance, 0.0, "Balance should be negative after overdraft")
|
||||
|
||||
s.Require().NoError(s.repo.UpdateConcurrency(s.ctx, user1.ID, 3), "UpdateConcurrency")
|
||||
got5, err := s.repo.GetByID(s.ctx, user1.ID)
|
||||
@@ -511,6 +532,6 @@ func (s *UserRepoSuite) TestUpdateConcurrency_NotFound() {
|
||||
func (s *UserRepoSuite) TestDeductBalance_NotFound() {
|
||||
err := s.repo.DeductBalance(s.ctx, 999999, 5)
|
||||
s.Require().Error(err, "expected error for non-existent user")
|
||||
// DeductBalance 在用户不存在时返回 ErrInsufficientBalance 因为 WHERE 条件不匹配
|
||||
s.Require().ErrorIs(err, service.ErrInsufficientBalance)
|
||||
// DeductBalance 在用户不存在时返回 ErrUserNotFound
|
||||
s.Require().ErrorIs(err, service.ErrUserNotFound)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user