From bab6c7bbf9a3942fa89ed399bc6c5f2971bedb15 Mon Sep 17 00:00:00 2001
From: zj <1772600164@qq.com>
Date: Sat, 30 Aug 2025 22:16:31 +0800
Subject: [PATCH] 1

---
 ruoyi-admin/src/main/java/com/ruoyi/im/service/impl/ImApiServcieImpl.java |   51 +++++++++++++++++++++++++++++++++++++++++++++++++--
 1 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/ruoyi-admin/src/main/java/com/ruoyi/im/service/impl/ImApiServcieImpl.java b/ruoyi-admin/src/main/java/com/ruoyi/im/service/impl/ImApiServcieImpl.java
index ada2ea4..41be005 100644
--- a/ruoyi-admin/src/main/java/com/ruoyi/im/service/impl/ImApiServcieImpl.java
+++ b/ruoyi-admin/src/main/java/com/ruoyi/im/service/impl/ImApiServcieImpl.java
@@ -121,6 +121,18 @@
             return Result.error("账号已被注册!");
         }
 
+        if(dto.getAccountType() == 0 && StringUtils.isEmpty(dto.getInvitationCode())){
+            return Result.error("邀请码不能为空!");
+        }
+
+        UserAccount user = userAccountService.getOne(new LambdaQueryWrapper<UserAccount>()
+                .eq(UserAccount::getInvitationCode, dto.getInvitationCode()).last(" limit 1"));
+        if(ObjectUtil.isEmpty(user)){
+            return Result.error("邀请码错误");
+        }
+
+        String invitationCode = getInvitationCode();
+
         // 创建本地用户账户记录
         UserAccount userAccount = new UserAccount();
         userAccount.setAccount(dto.getAccount());
@@ -131,6 +143,8 @@
         userAccount.setNickname(dto.getNickname());
         userAccount.setCreateTime(new Date());
         userAccount.setUpdateTime(new Date());
+        userAccount.setInvitationCode(invitationCode);
+        userAccount.setInvitationAccount(user.getAccount());
 
         if (!userAccountService.save(userAccount)) {
             throw new RuntimeException("保存用户账户失败");
@@ -184,6 +198,38 @@
         }
     }
 
+    private String getInvitationCode() {
+        String invitationCode = null;
+        int maxAttempts = 100; // 最大尝试次数
+        int attempts = 0;
+
+        while (attempts < maxAttempts) {
+            invitationCode = generateInvitationCode();
+            long count = userAccountService.count(new LambdaQueryWrapper<UserAccount>()
+                    .eq(UserAccount::getInvitationCode, invitationCode));
+            if(count <= 0){
+                break;
+            }
+            attempts++;
+        }
+
+        if (attempts >= maxAttempts) {
+            log.error("生成邀请码已超最大尝试次数!");
+            throw new RuntimeException("无法生成唯一的邀请码,请稍后重试");
+        }
+        return invitationCode;
+    }
+
+
+    /**
+     * 生成邀请码
+     * @return
+     */
+    public static String generateInvitationCode() {
+        Random random = new Random();
+        int code = 100000 + random.nextInt(900000);
+        return String.valueOf(code);
+    }
 
     /**
      * 优化的账号生成方法,使用雪花算法变体提高并发性能
@@ -440,6 +486,7 @@
     @Override
     @Transactional(rollbackFor = Exception.class)
     public Result batchRegister(RegisterDto dto) {
+        dto.setAccountType(1);
         if(dto.getType() == 2){
             return register(dto);
         }else{
@@ -492,7 +539,7 @@
         List<UserAccount> accounts = new ArrayList<>(count);
         Set<String> generatedAccounts = new HashSet<>(count); // 用于内存中去重
         Random random = new Random();
-
+        String invitationCode = getInvitationCode();
         for (int i = 0; i < count; i++) {
             String account;
             do {
@@ -510,7 +557,7 @@
             userAccount.setCreateTime(new Date());
             userAccount.setNickname("用户_" + account.substring(7)); // 简单生成昵称
             userAccount.setAccountType(1); // 设置账号类型为1
-
+            userAccount.setInvitationCode(invitationCode);
             accounts.add(userAccount);
         }
         return accounts;

--
Gitblit v1.9.3