router.go 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150
  1. package router
  2. import (
  3. "github.com/gogf/gf/v2/errors/gcode"
  4. "github.com/gogf/gf/v2/errors/gerror"
  5. "github.com/gogf/gf/v2/net/ghttp"
  6. "github.com/gogf/gf/v2/os/glog"
  7. "go-gpt/internal/controller"
  8. "go-gpt/internal/service"
  9. "net/http"
  10. "sync"
  11. "time"
  12. )
  13. type DefaultHandlerResponse struct {
  14. Code int `json:"status"`
  15. Message string `json:"message"`
  16. Data interface{} `json:"data"`
  17. }
  18. // MiddlewareHandlerResponse is the default middleware handling handler response object and its error.
  19. func MiddlewareHandlerResponse(r *ghttp.Request) {
  20. r.Middleware.Next()
  21. // There's custom buffer content, it then exits current handler.
  22. if r.Response.BufferLength() > 0 {
  23. return
  24. }
  25. var (
  26. msg string
  27. err = r.GetError()
  28. res = r.GetHandlerResponse()
  29. code = gerror.Code(err)
  30. )
  31. if err != nil {
  32. if code == gcode.CodeNil {
  33. code = gcode.CodeInternalError
  34. }
  35. msg = err.Error()
  36. } else {
  37. if r.Response.Status > 0 && r.Response.Status != http.StatusOK {
  38. msg = http.StatusText(r.Response.Status)
  39. switch r.Response.Status {
  40. case http.StatusNotFound:
  41. code = gcode.CodeNotFound
  42. case http.StatusForbidden:
  43. code = gcode.CodeNotAuthorized
  44. default:
  45. code = gcode.CodeUnknown
  46. }
  47. // It creates error as it can be retrieved by other middlewares.
  48. err = gerror.NewCode(code, msg)
  49. r.SetError(err)
  50. } else {
  51. code = gcode.CodeOK
  52. }
  53. }
  54. r.Response.WriteJson(DefaultHandlerResponse{
  55. Code: code.Code(),
  56. Message: msg,
  57. Data: res,
  58. })
  59. }
  60. type middlewareService struct{}
  61. func NewTokenMiddleware() *middlewareService {
  62. return &middlewareService{}
  63. }
  64. func (s *middlewareService) Auth(r *ghttp.Request) {
  65. service.Auth().MiddlewareFunc()(r)
  66. r.Middleware.Next()
  67. }
  68. type RateLimitMiddleware struct {
  69. sync.Mutex
  70. Counter map[string]int64
  71. }
  72. func NewRateLimitMiddleware() *RateLimitMiddleware {
  73. return &RateLimitMiddleware{Counter: make(map[string]int64)}
  74. }
  75. func (m *RateLimitMiddleware) Middleware(r *ghttp.Request) {
  76. ip := r.GetClientIp()
  77. m.Lock()
  78. defer m.Unlock()
  79. now := time.Now().Unix()
  80. count, ok := m.Counter[ip]
  81. if !ok {
  82. count = 1
  83. m.Counter[ip] = now
  84. } else {
  85. if now-m.Counter[ip] < 1 {
  86. r.Response.WriteStatusExit(429, "Too Many Requests")
  87. } else {
  88. m.Counter[ip] = now
  89. count++
  90. }
  91. }
  92. glog.Debugf(r.Context(), "%s:%d", ip, count)
  93. r.Middleware.Next()
  94. }
  95. func MiddlewareCORS(r *ghttp.Request) {
  96. corsOptions := r.Response.DefaultCORSOptions()
  97. //corsOptions.AllowDomain = []string{"goframe.org", "baidu.com"}
  98. r.Response.CORS(corsOptions)
  99. r.Response.CORSDefault()
  100. r.Middleware.Next()
  101. }
  102. func BindController(group *ghttp.RouterGroup) {
  103. DomeRouter(group)
  104. group.Group("/api/v1", func(group *ghttp.RouterGroup) {
  105. vv := NewRateLimitMiddleware()
  106. group.Middleware(MiddlewareCORS, vv.Middleware)
  107. SessionRouter(group)
  108. ChatRouter(group)
  109. })
  110. }
  111. func DomeRouter(group *ghttp.RouterGroup) {
  112. //tokenMi := NewTokenMiddleware()
  113. //group.Middleware(tokenMi.Auth)
  114. group.Bind(
  115. controller.Hello,
  116. )
  117. }
  118. func ChatRouter(group *ghttp.RouterGroup) {
  119. //group.Middleware(MiddlewareHandlerResponse)
  120. group.Bind(
  121. controller.Chat,
  122. )
  123. }
  124. func SessionRouter(group *ghttp.RouterGroup) {
  125. //group.Middleware(MiddlewareHandlerResponse)
  126. group.Group("/session", func(group *ghttp.RouterGroup) {
  127. group.Bind(
  128. controller.Session,
  129. )
  130. })
  131. }