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

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