25개 이상의 토픽을 선택하실 수 없습니다. Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

157 lines
4.5 KiB

  1. package region_sync
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "fmt"
  6. "io"
  7. "net/http"
  8. "strings"
  9. "time"
  10. "github.com/QuantumNous/new-api/common"
  11. "github.com/QuantumNous/new-api/logger"
  12. )
  13. // SyncClient 调用远程节点的 HTTP 客户端
  14. type SyncClient struct {
  15. endpoint string
  16. apiKey string
  17. httpClient *http.Client
  18. }
  19. func NewSyncClient(endpoint, apiKey string) *SyncClient {
  20. return &SyncClient{
  21. endpoint: endpoint,
  22. apiKey: apiKey,
  23. httpClient: &http.Client{Timeout: 30 * time.Second},
  24. }
  25. }
  26. func (c *SyncClient) doRequest(method, path string, body interface{}) ([]byte, error) {
  27. var reqBody io.Reader
  28. if body != nil {
  29. data, err := json.Marshal(body)
  30. if err != nil {
  31. return nil, fmt.Errorf("marshal request: %w", err)
  32. }
  33. reqBody = bytes.NewReader(data)
  34. }
  35. // 处理 URL 拼接,避免双斜杠问题
  36. endpoint := strings.TrimRight(c.endpoint, "/")
  37. path = strings.TrimLeft(path, "/")
  38. url := endpoint + "/" + path
  39. common.SysLog(fmt.Sprintf("[RegionSync] HTTP %s %s", method, url))
  40. logger.LogDebug(nil, "[RegionSync] HTTP %s %s", method, url)
  41. req, err := http.NewRequest(method, url, reqBody)
  42. if err != nil {
  43. return nil, fmt.Errorf("create request: %w", err)
  44. }
  45. req.Header.Set("Content-Type", "application/json")
  46. req.Header.Set("X-Sync-API-Key", c.apiKey)
  47. req.Header.Set("X-Sync-Node", c.endpoint)
  48. resp, err := c.httpClient.Do(req)
  49. if err != nil {
  50. return nil, fmt.Errorf("http request: %w", err)
  51. }
  52. defer resp.Body.Close()
  53. data, err := io.ReadAll(resp.Body)
  54. if err != nil {
  55. return nil, fmt.Errorf("read response: %w", err)
  56. }
  57. // 如果返回状态码不是 200,记录错误信息
  58. if resp.StatusCode != http.StatusOK {
  59. common.SysError(fmt.Sprintf("[RegionSync] HTTP error %d from %s: %s", resp.StatusCode, url, string(data)))
  60. return nil, fmt.Errorf("http status %d: %s", resp.StatusCode, string(data))
  61. }
  62. return data, nil
  63. }
  64. func (c *SyncClient) SyncUserCreate(req *SyncUserRequest) (*SyncUserResponse, error) {
  65. data, err := c.doRequest("POST", "/api/internal/sync/user/create", req)
  66. if err != nil {
  67. return nil, err
  68. }
  69. var resp SyncUserResponse
  70. if err := json.Unmarshal(data, &resp); err != nil {
  71. return nil, fmt.Errorf("unmarshal response: %w", err)
  72. }
  73. return &resp, nil
  74. }
  75. func (c *SyncClient) QueryQuota(userId int) (*QueryQuotaResponse, error) {
  76. data, err := c.doRequest("POST", "/api/internal/sync/quota/query", &QueryQuotaRequest{UserId: userId})
  77. if err != nil {
  78. return nil, err
  79. }
  80. var resp QueryQuotaResponse
  81. if err := json.Unmarshal(data, &resp); err != nil {
  82. return nil, fmt.Errorf("unmarshal response: %w", err)
  83. }
  84. logger.LogDebug(nil, "[RegionSync] QueryQuota: userId=%d, quota=%d, success=%v", userId, resp.Quota, resp.Success)
  85. return &resp, nil
  86. }
  87. func (c *SyncClient) BatchDeduct(req *BatchDeductRequest) (*BatchDeductResponse, error) {
  88. logger.LogDebug(nil, "[RegionSync] BatchDeduct: sending %d records", len(req.Records))
  89. data, err := c.doRequest("POST", "/api/internal/sync/quota/batch-deduct", req)
  90. if err != nil {
  91. return nil, err
  92. }
  93. var resp BatchDeductResponse
  94. if err := json.Unmarshal(data, &resp); err != nil {
  95. return nil, fmt.Errorf("unmarshal response: %w", err)
  96. }
  97. successCount := 0
  98. for _, r := range resp.Results {
  99. if r.Success {
  100. successCount++
  101. }
  102. }
  103. logger.LogDebug(nil, "[RegionSync] BatchDeduct: response success=%d, total=%d", successCount, len(resp.Results))
  104. return &resp, nil
  105. }
  106. func (c *SyncClient) UpdateQuota(remoteUserId, quota int) (*UpdateQuotaResponse, error) {
  107. data, err := c.doRequest("POST", "/api/internal/sync/quota/update", &UpdateQuotaRequest{RemoteUserId: remoteUserId, Quota: quota})
  108. if err != nil {
  109. return nil, err
  110. }
  111. var resp UpdateQuotaResponse
  112. if err := json.Unmarshal(data, &resp); err != nil {
  113. return nil, fmt.Errorf("unmarshal response: %w", err)
  114. }
  115. return &resp, nil
  116. }
  117. func (c *SyncClient) FetchConfig() (*SyncConfigResponse, error) {
  118. data, err := c.doRequest("GET", "/api/internal/sync/config", nil)
  119. if err != nil {
  120. return nil, err
  121. }
  122. var resp SyncConfigResponse
  123. if err := json.Unmarshal(data, &resp); err != nil {
  124. return nil, fmt.Errorf("unmarshal response: %w", err)
  125. }
  126. return &resp, nil
  127. }
  128. func (c *SyncClient) BatchQueryQuota(userIds []int) (*BatchQueryQuotaResponse, error) {
  129. data, err := c.doRequest("POST", "/api/internal/sync/quota/batch-query", &BatchQueryQuotaRequest{UserIds: userIds})
  130. if err != nil {
  131. return nil, err
  132. }
  133. var resp BatchQueryQuotaResponse
  134. if err := json.Unmarshal(data, &resp); err != nil {
  135. return nil, fmt.Errorf("unmarshal response: %w", err)
  136. }
  137. return &resp, nil
  138. }