middleware.go 2.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283
  1. package router
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "net/http"
  7. "time"
  8. "github.com/mhaya/game/game_cluster/internal/mdb/models"
  9. "go.mongodb.org/mongo-driver/bson"
  10. "github.com/gin-gonic/gin"
  11. "github.com/mhaya/game/game_cluster/internal/constant"
  12. "github.com/mhaya/game/game_cluster/internal/mdb"
  13. )
  14. func Auth() gin.HandlerFunc {
  15. return func(c *gin.Context) {
  16. tokenString := c.GetHeader("Token")
  17. if tokenString == "" {
  18. c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
  19. "msg": "token is empty",
  20. })
  21. return
  22. }
  23. result, err := mdb.RDB.Get(context.Background(), tokenString).Result()
  24. if err != nil {
  25. return
  26. }
  27. if result == "" {
  28. c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
  29. "msg": "token is invalid",
  30. })
  31. return
  32. }
  33. // 获取请求URL
  34. url := c.Request.URL.Path
  35. s, _ := mdb.RDB.HGet(context.Background(), "admin::token::"+tokenString, url).Result()
  36. admin, _ := mdb.RDB.HGet(context.Background(), "admin::token::"+tokenString, constant.AdminAccess).Result()
  37. // 检查是否有权限
  38. if s == "" && admin == "" {
  39. c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
  40. "msg": "token is no auth",
  41. })
  42. return
  43. }
  44. if admin == "" {
  45. if err := checkIP(c); err != nil {
  46. c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
  47. "msg": "ip is no auth",
  48. })
  49. return
  50. }
  51. }
  52. c.Next()
  53. }
  54. }
  55. // checkIP
  56. func checkIP(c *gin.Context) error {
  57. // 获取请求的ip
  58. ip := c.ClientIP()
  59. whitelistModel := &models.Whitelist{}
  60. collection := mdb.MDB.Collection(whitelistModel.TableName())
  61. // 设置超时时间
  62. ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
  63. defer cancel() // 确保在函数退出时取消上下文
  64. // 示例:查询 IP 是否在白名单中
  65. err := collection.FindOne(ctx, bson.M{"ip": ip}).Decode(&whitelistModel)
  66. if err != nil {
  67. // 处理查询错误
  68. return fmt.Errorf("failed to check IP in whitelist: %w", err)
  69. }
  70. // 根据查询结果决定是否允许访问
  71. if whitelistModel != nil {
  72. return nil // 允许访问
  73. } else {
  74. return errors.New("IP not in whitelist") // 拒绝访问
  75. }
  76. }