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.
 
 
 

195 lines
5.5 KiB

  1. package router
  2. import (
  3. "net/http"
  4. "net/http/httptest"
  5. "testing"
  6. "github.com/QuantumNous/new-api/common"
  7. "github.com/QuantumNous/new-api/controller"
  8. "github.com/QuantumNous/new-api/middleware"
  9. "github.com/QuantumNous/new-api/model"
  10. "github.com/gin-contrib/sessions"
  11. "github.com/gin-contrib/sessions/cookie"
  12. "github.com/gin-gonic/gin"
  13. "github.com/glebarez/sqlite"
  14. "github.com/stretchr/testify/require"
  15. "gorm.io/gorm"
  16. )
  17. func setupUserMigrationRouterDB(t *testing.T) *gorm.DB {
  18. t.Helper()
  19. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  20. require.NoError(t, err)
  21. sqlDB, err := db.DB()
  22. require.NoError(t, err)
  23. sqlDB.SetMaxOpenConns(1)
  24. require.NoError(t, db.AutoMigrate(
  25. &model.User{},
  26. &model.CustomOAuthProvider{},
  27. &model.UserOAuthBinding{},
  28. &model.UserMigrationBatch{},
  29. &model.UserMigrationItem{},
  30. &model.MigrationQuotaGrant{},
  31. ))
  32. return db
  33. }
  34. func setupUserMigrationRootRouter() *gin.Engine {
  35. gin.SetMode(gin.TestMode)
  36. r := gin.New()
  37. store := cookie.NewStore([]byte("test-session-secret"))
  38. store.Options(sessions.Options{
  39. Path: "/",
  40. MaxAge: 3600,
  41. HttpOnly: true,
  42. Secure: false,
  43. })
  44. r.Use(sessions.Sessions("session", store))
  45. userMigrationRoute := r.Group("/api/user-migrations")
  46. userMigrationRoute.Use(func(c *gin.Context) {
  47. session := sessions.Default(c)
  48. session.Set("username", "root")
  49. session.Set("role", common.RoleRootUser)
  50. session.Set("id", 1)
  51. session.Set("status", common.UserStatusEnabled)
  52. session.Save()
  53. c.Request.Header.Set("New-Api-User", "1")
  54. c.Next()
  55. }, middleware.RootAuth(), middleware.MasterNodeOnly())
  56. {
  57. userMigrationRoute.GET("/batches", controller.ListUserMigrationBatches)
  58. userMigrationRoute.POST("/batches/:id/execute", controller.ExecuteUserMigrationBatch)
  59. userMigrationRoute.POST("/batches/:id/cancel", controller.CancelUserMigrationBatch)
  60. }
  61. return r
  62. }
  63. func setupUserMigrationAdminRouter() *gin.Engine {
  64. gin.SetMode(gin.TestMode)
  65. r := gin.New()
  66. store := cookie.NewStore([]byte("test-session-secret"))
  67. store.Options(sessions.Options{
  68. Path: "/",
  69. MaxAge: 3600,
  70. HttpOnly: true,
  71. Secure: false,
  72. })
  73. r.Use(sessions.Sessions("session", store))
  74. userMigrationRoute := r.Group("/api/user-migrations")
  75. userMigrationRoute.Use(func(c *gin.Context) {
  76. session := sessions.Default(c)
  77. session.Set("username", "admin")
  78. session.Set("role", common.RoleAdminUser)
  79. session.Set("id", 2)
  80. session.Set("status", common.UserStatusEnabled)
  81. session.Save()
  82. c.Request.Header.Set("New-Api-User", "2")
  83. c.Next()
  84. }, middleware.RootAuth(), middleware.MasterNodeOnly())
  85. {
  86. userMigrationRoute.GET("/batches", controller.ListUserMigrationBatches)
  87. userMigrationRoute.POST("/batches/:id/execute", controller.ExecuteUserMigrationBatch)
  88. userMigrationRoute.POST("/batches/:id/cancel", controller.CancelUserMigrationBatch)
  89. }
  90. return r
  91. }
  92. func TestUserMigrationRoutes_RootCanAccess(t *testing.T) {
  93. db := setupUserMigrationRouterDB(t)
  94. orig := model.DB
  95. origMaster := common.IsMasterNode
  96. model.DB = db
  97. defer func() {
  98. model.DB = orig
  99. common.IsMasterNode = origMaster
  100. }()
  101. common.IsMasterNode = true
  102. require.NoError(t, db.Create(&model.UserMigrationBatch{
  103. Id: 1,
  104. Name: "router-batch",
  105. SourceRegion: "ov",
  106. TargetRegion: "cn",
  107. OperatorId: 1,
  108. Status: model.UserMigrationBatchStatusDraft,
  109. }).Error)
  110. router := setupUserMigrationRootRouter()
  111. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches", nil)
  112. w := httptest.NewRecorder()
  113. router.ServeHTTP(w, req)
  114. require.Equal(t, http.StatusOK, w.Code)
  115. req = httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil)
  116. w = httptest.NewRecorder()
  117. router.ServeHTTP(w, req)
  118. require.NotEqual(t, http.StatusUnauthorized, w.Code)
  119. req = httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/cancel", nil)
  120. w = httptest.NewRecorder()
  121. router.ServeHTTP(w, req)
  122. require.NotEqual(t, http.StatusUnauthorized, w.Code)
  123. }
  124. func TestUserMigrationRoutes_AdminForbidden(t *testing.T) {
  125. db := setupUserMigrationRouterDB(t)
  126. orig := model.DB
  127. origMaster := common.IsMasterNode
  128. model.DB = db
  129. defer func() {
  130. model.DB = orig
  131. common.IsMasterNode = origMaster
  132. }()
  133. common.IsMasterNode = true
  134. router := setupUserMigrationAdminRouter()
  135. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches", nil)
  136. w := httptest.NewRecorder()
  137. router.ServeHTTP(w, req)
  138. require.Equal(t, http.StatusOK, w.Code)
  139. require.Contains(t, w.Body.String(), "success\":false")
  140. req = httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/execute", nil)
  141. w = httptest.NewRecorder()
  142. router.ServeHTTP(w, req)
  143. require.Equal(t, http.StatusOK, w.Code)
  144. require.Contains(t, w.Body.String(), "success\":false")
  145. req = httptest.NewRequest(http.MethodPost, "/api/user-migrations/batches/1/cancel", nil)
  146. w = httptest.NewRecorder()
  147. router.ServeHTTP(w, req)
  148. require.Equal(t, http.StatusOK, w.Code)
  149. require.Contains(t, w.Body.String(), "success\":false")
  150. }
  151. func TestUserMigrationRoutes_RootForbiddenOnSlaveNode(t *testing.T) {
  152. db := setupUserMigrationRouterDB(t)
  153. orig := model.DB
  154. origMaster := common.IsMasterNode
  155. model.DB = db
  156. defer func() {
  157. model.DB = orig
  158. common.IsMasterNode = origMaster
  159. }()
  160. common.IsMasterNode = false
  161. router := setupUserMigrationRootRouter()
  162. req := httptest.NewRequest(http.MethodGet, "/api/user-migrations/batches", nil)
  163. w := httptest.NewRecorder()
  164. router.ServeHTTP(w, req)
  165. require.Equal(t, http.StatusForbidden, w.Code)
  166. require.Contains(t, w.Body.String(), "only master node")
  167. }