Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

277 строки
7.8 KiB

  1. package user_migration
  2. import (
  3. "fmt"
  4. "strings"
  5. "github.com/QuantumNous/new-api/common"
  6. "github.com/QuantumNous/new-api/model"
  7. "gorm.io/gorm"
  8. )
  9. func (s *Service) ExecuteBatch(batchId int) error {
  10. items, err := model.ListExecutableUserMigrationItems(batchId)
  11. if err != nil {
  12. return err
  13. }
  14. for _, item := range items {
  15. if item.Status == model.UserMigrationItemStatusMigrated {
  16. continue
  17. }
  18. if err := s.executeItem(item); err != nil {
  19. _ = model.MarkUserMigrationItemFailed(item.Id, err.Error())
  20. }
  21. }
  22. return model.RefreshUserMigrationBatchExecutionStats(batchId)
  23. }
  24. func (s *Service) executeItem(item *model.UserMigrationItem) error {
  25. switch item.ResolutionStrategy {
  26. case model.UserMigrationStrategyCreateNew:
  27. return s.executeCreateNew(item)
  28. case model.UserMigrationStrategyMergeExisting:
  29. return s.executeMergeExisting(item)
  30. case model.UserMigrationStrategySkip:
  31. return model.MarkUserMigrationItemSkipped(item.Id)
  32. default:
  33. return fmt.Errorf("unsupported strategy: %q", item.ResolutionStrategy)
  34. }
  35. }
  36. func (s *Service) executeCreateNew(item *model.UserMigrationItem) error {
  37. if item.TargetUserId != 0 {
  38. cnUser, err := model.GetUserById(item.TargetUserId, true)
  39. if err != nil {
  40. return err
  41. }
  42. return s.ensureRemoteConverted(item, cnUser.Quota)
  43. }
  44. snapshot, err := s.client.GetUser(item.SourceUserId)
  45. if err != nil {
  46. return fmt.Errorf("get remote user %d: %w", item.SourceUserId, err)
  47. }
  48. if err := s.revalidateCreateNew(item, snapshot); err != nil {
  49. return err
  50. }
  51. if err := markItemMigrating(item.Id); err != nil {
  52. return err
  53. }
  54. affCode := snapshot.AffCode
  55. if affCode != "" {
  56. if _, err := model.GetUserIdByAffCode(affCode); err == nil {
  57. affCode = common.GetRandomString(4)
  58. }
  59. }
  60. cnUser := &model.User{
  61. Username: snapshot.Username,
  62. Password: snapshot.Password,
  63. Email: snapshot.Email,
  64. DisplayName: snapshot.DisplayName,
  65. Status: snapshot.Status,
  66. Role: snapshot.Role,
  67. Group: snapshot.Group,
  68. Quota: snapshot.Quota,
  69. AffCode: affCode,
  70. CreatedAt: snapshot.CreatedAt,
  71. Setting: snapshot.Setting,
  72. GitHubId: snapshot.GitHubId,
  73. DiscordId: snapshot.DiscordId,
  74. OidcId: snapshot.OidcId,
  75. WeChatId: snapshot.WeChatId,
  76. TelegramId: snapshot.TelegramId,
  77. LinuxDOId: snapshot.LinuxDOId,
  78. }
  79. err = model.DB.Transaction(func(tx *gorm.DB) error {
  80. if err := model.InsertImportedUserWithHashedPasswordTx(tx, cnUser); err != nil {
  81. if isUniqueConstraintErr(err) {
  82. return fmt.Errorf("execution-time identity drift detected, rescan required: %w", err)
  83. }
  84. return fmt.Errorf("create CN user: %w", err)
  85. }
  86. if err := copyMappedBindingsTx(tx, cnUser.Id, snapshot.OAuthBindings); err != nil {
  87. return err
  88. }
  89. if err := tx.Model(&model.UserMigrationItem{}).Where("id = ?", item.Id).
  90. Update("target_user_id", cnUser.Id).Error; err != nil {
  91. return err
  92. }
  93. item.TargetUserId = cnUser.Id
  94. return nil
  95. })
  96. if err != nil {
  97. return err
  98. }
  99. return s.ensureRemoteConverted(item, cnUser.Quota)
  100. }
  101. func (s *Service) executeMergeExisting(item *model.UserMigrationItem) error {
  102. snapshot, err := s.client.GetUser(item.SourceUserId)
  103. if err != nil {
  104. return fmt.Errorf("get remote user %d: %w", item.SourceUserId, err)
  105. }
  106. if err := s.revalidateMergeExisting(item, snapshot); err != nil {
  107. return err
  108. }
  109. if err := markItemMigrating(item.Id); err != nil {
  110. return err
  111. }
  112. finalQuota := 0
  113. err = model.DB.Transaction(func(tx *gorm.DB) error {
  114. _, err := model.AtomicGrantQuotaIfNotExists(tx, item.TargetUserId, snapshot.Quota, item.MigrationToken)
  115. if err != nil {
  116. return err
  117. }
  118. var cnUser model.User
  119. if err := tx.Select("quota").First(&cnUser, item.TargetUserId).Error; err != nil {
  120. return err
  121. }
  122. finalQuota = cnUser.Quota
  123. return copyMappedBindingsTx(tx, item.TargetUserId, snapshot.OAuthBindings)
  124. })
  125. if err != nil {
  126. return err
  127. }
  128. return s.ensureRemoteConverted(item, finalQuota)
  129. }
  130. func copyMappedBindingsTx(tx *gorm.DB, userId int, bindings []RemoteOAuthBinding) error {
  131. if len(bindings) == 0 {
  132. return nil
  133. }
  134. slugs := make([]string, 0, len(bindings))
  135. for _, binding := range bindings {
  136. slugs = append(slugs, binding.ProviderSlug)
  137. }
  138. providers, err := model.GetCustomOAuthProvidersBySlugs(slugs)
  139. if err != nil {
  140. return err
  141. }
  142. slugToID := make(map[string]int, len(providers))
  143. for _, provider := range providers {
  144. slugToID[provider.Slug] = provider.Id
  145. }
  146. localBindings := make([]*model.UserOAuthBinding, 0, len(bindings))
  147. for _, binding := range bindings {
  148. localID, ok := slugToID[binding.ProviderSlug]
  149. if !ok {
  150. continue
  151. }
  152. localBindings = append(localBindings, &model.UserOAuthBinding{
  153. ProviderId: localID,
  154. ProviderUserId: binding.ProviderUserId,
  155. })
  156. }
  157. return model.CopyOAuthBindingsIfMissingTx(tx, userId, localBindings)
  158. }
  159. func (s *Service) ensureRemoteConverted(item *model.UserMigrationItem, syncedQuota int) error {
  160. remoteUser, err := s.client.GetUser(item.SourceUserId)
  161. if err != nil {
  162. return err
  163. }
  164. if remoteUser.Source != common.UserSourceSynced || remoteUser.RemoteUserId != item.TargetUserId {
  165. if err := s.client.ConvertToSynced(item.SourceUserId, &ConvertRemoteUserRequest{
  166. RemoteUserId: item.TargetUserId,
  167. SyncedQuota: syncedQuota,
  168. }); err != nil {
  169. return err
  170. }
  171. }
  172. return markItemMigrated(item.Id)
  173. }
  174. func (s *Service) revalidateCreateNew(item *model.UserMigrationItem, snapshot *RemoteUserSnapshot) error {
  175. if snapshot.Source == common.UserSourceSynced {
  176. return fmt.Errorf("execution-time identity drift detected, remote user %d is already synced, rescan required", item.SourceUserId)
  177. }
  178. result, err := s.analyzeExecutionSnapshot(snapshot)
  179. if err != nil {
  180. return err
  181. }
  182. if result.Status != model.UserMigrationItemStatusReady ||
  183. result.Strategy != model.UserMigrationStrategyCreateNew ||
  184. result.TargetCNUserID != 0 {
  185. return fmt.Errorf("execution-time identity drift detected, rescan required")
  186. }
  187. return nil
  188. }
  189. func (s *Service) revalidateMergeExisting(item *model.UserMigrationItem, snapshot *RemoteUserSnapshot) error {
  190. if snapshot.Source == common.UserSourceSynced {
  191. if snapshot.RemoteUserId == item.TargetUserId {
  192. return nil
  193. }
  194. return fmt.Errorf("execution-time identity drift detected, remote user %d is already synced to CN user %d", item.SourceUserId, snapshot.RemoteUserId)
  195. }
  196. result, err := s.analyzeExecutionSnapshot(snapshot)
  197. if err != nil {
  198. return err
  199. }
  200. if result.Status != model.UserMigrationItemStatusReady ||
  201. result.Strategy != model.UserMigrationStrategyMergeExisting ||
  202. result.TargetCNUserID != item.TargetUserId {
  203. return fmt.Errorf("execution-time merge target drift detected, rescan required")
  204. }
  205. return nil
  206. }
  207. func (s *Service) analyzeExecutionSnapshot(snapshot *RemoteUserSnapshot) (*conflictResult, error) {
  208. providerSlugMap, err := buildProviderSlugMap()
  209. if err != nil {
  210. return nil, err
  211. }
  212. result, err := analyzeConflicts(snapshot, providerSlugMap, s.client)
  213. if err != nil {
  214. return nil, fmt.Errorf("revalidate conflicts for user %d: %w", snapshot.Id, err)
  215. }
  216. return result, nil
  217. }
  218. func markItemMigrating(itemID int) error {
  219. return model.DB.Model(&model.UserMigrationItem{}).Where("id = ?", itemID).Updates(map[string]any{
  220. "status": model.UserMigrationItemStatusMigrating,
  221. "error_message": "",
  222. }).Error
  223. }
  224. func markItemMigrated(itemID int) error {
  225. return model.DB.Model(&model.UserMigrationItem{}).Where("id = ?", itemID).Updates(map[string]any{
  226. "status": model.UserMigrationItemStatusMigrated,
  227. "migrated_at": timeNow(),
  228. "error_message": "",
  229. }).Error
  230. }
  231. func isUniqueConstraintErr(err error) bool {
  232. if err == nil {
  233. return false
  234. }
  235. message := err.Error()
  236. return strings.Contains(message, "UNIQUE constraint failed") ||
  237. strings.Contains(message, "Duplicate entry") ||
  238. strings.Contains(strings.ToLower(message), "duplicate key value")
  239. }
  240. func timeNow() int64 {
  241. return common.GetTimestamp()
  242. }