You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

83 lines
1.7 KiB

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