Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.
 
 
 

94 řádky
2.1 KiB

  1. package common
  2. import (
  3. "io"
  4. "net/http"
  5. "strings"
  6. commonpkg "github.com/QuantumNous/new-api/common"
  7. "github.com/gin-gonic/gin"
  8. )
  9. const (
  10. ContextKeyRelayChatID = "relay_chat_id"
  11. ContextKeyRelayUpstreamID = "relay_upstream_id"
  12. )
  13. func SetRelayChatID(c *gin.Context, chatID string) {
  14. if c == nil {
  15. return
  16. }
  17. chatID = strings.TrimSpace(chatID)
  18. if chatID == "" {
  19. return
  20. }
  21. c.Set(ContextKeyRelayChatID, chatID)
  22. }
  23. func GetRelayChatID(c *gin.Context) string {
  24. if c == nil {
  25. return ""
  26. }
  27. return strings.TrimSpace(c.GetString(ContextKeyRelayChatID))
  28. }
  29. func SetRelayUpstreamID(c *gin.Context, upstreamID string) {
  30. if c == nil {
  31. return
  32. }
  33. upstreamID = strings.TrimSpace(upstreamID)
  34. if upstreamID == "" {
  35. return
  36. }
  37. c.Set(ContextKeyRelayUpstreamID, upstreamID)
  38. }
  39. func GetRelayUpstreamID(c *gin.Context) string {
  40. if c == nil {
  41. return ""
  42. }
  43. return strings.TrimSpace(c.GetString(ContextKeyRelayUpstreamID))
  44. }
  45. func ExtractTopLevelChatID(c *gin.Context) (string, error) {
  46. if c == nil {
  47. return "", nil
  48. }
  49. if !strings.HasPrefix(strings.ToLower(strings.TrimSpace(c.Request.Header.Get("Content-Type"))), "application/json") {
  50. return "", nil
  51. }
  52. bodyStorage, err := commonpkg.GetBodyStorage(c)
  53. if err != nil {
  54. return "", err
  55. }
  56. bodyBytes, err := bodyStorage.Bytes()
  57. if err != nil || len(bodyBytes) == 0 {
  58. if _, seekErr := bodyStorage.Seek(0, io.SeekStart); seekErr == nil {
  59. c.Request.Body = io.NopCloser(bodyStorage)
  60. }
  61. return "", err
  62. }
  63. if _, seekErr := bodyStorage.Seek(0, io.SeekStart); seekErr != nil {
  64. return "", seekErr
  65. }
  66. c.Request.Body = io.NopCloser(bodyStorage)
  67. var payload map[string]interface{}
  68. if err := commonpkg.Unmarshal(bodyBytes, &payload); err != nil {
  69. return "", nil
  70. }
  71. return strings.TrimSpace(commonpkg.Interface2String(payload["chat_id"])), nil
  72. }
  73. func CaptureUpstreamIDFromHTTPResponse(info *RelayInfo, resp *http.Response) {
  74. if info == nil || resp == nil {
  75. return
  76. }
  77. for _, key := range []string{"x-request-id", "request-id"} {
  78. value := strings.TrimSpace(resp.Header.Get(key))
  79. if value != "" {
  80. info.UpstreamID = value
  81. return
  82. }
  83. }
  84. }