邀请码功能

1、新增邀请码功能
  needInvitationCode: false 改成 true 后生效
2、完善后台用户注册 修改相关

pkg/common/config/config.go

@@ -492,6 +492,7 @@ type config struct {
 		}
 		TestDepartMentID     string   `yaml:"testDepartMentID"`
 		ImAPIURL             string   `yaml:"imAPIURL"`
+		NeedInvitationCode   bool     `yaml:"needInvitationCode"`
 		OnboardProcess       bool     `yaml:"onboardProcess"`
 		JoinDepartmentIDList []string `yaml:"joinDepartmentIDList"`
 		JoinDepartmentGroups bool     `yaml:"joinDepartmentGroups"`

pkg/common/constant/error.go

@@ -50,6 +50,7 @@ var (
 	DBMsg                       = errors.New("db failed")
 	ArgsMsg                     = errors.New("args failed")
 	CallBackMsg                 = errors.New("callback failed")
+	InvitationMsg               = errors.New("invitationCode error")
 
 	ThirdPartyMsg = errors.New("third party error")
 )
@@ -69,6 +70,7 @@ const (
 	ResetPasswordFailed  = 10011
 	RegisterLimit        = 10012
 	LoginLimit           = 10013
+	InvitationError      = 10014
 	DatabaseError        = 10002
 	ServerError          = 10004
 	HttpError            = 10005

pkg/common/db/model_struct.go

@@ -12,8 +12,11 @@ type Register struct {
 }
 
 type Invitation struct {
-	InvitationCode string    `gorm:"column:invitation_code;primary_key;type:varchar(255)"`
+	InvitationCode string    `gorm:"column:invitation_code;primary_key;type:varchar(32)"`
 	CreateTime     time.Time `gorm:"column:create_time"`
+	UserID         string    `gorm:"column:user_id"`
+	LastTime       time.Time `gorm:"column:last_time"`
+	Status         int32     `gorm:"column:status"`
 }
 
 //
@@ -180,6 +183,7 @@ type User struct {
 	LoginLimit       int32     `gorm:"column:login_limit"`
 	AppMangerLevel   int32     `gorm:"column:app_manger_level"`
 	GlobalRecvMsgOpt int32     `gorm:"column:global_recv_msg_opt"`
+	InvitationCode   string    `gorm:"column:invitation_code"`
 	status           int32     `gorm:"column:status"`
 }
 

pkg/common/db/mysql.go

@@ -159,6 +159,10 @@ func initMysqlDB() {
 		fmt.Println("CreateTable RegisterAddFriend")
 		db.Migrator().CreateTable(&RegisterAddFriend{})
 	}
+	if !db.Migrator().HasTable(&Invitation{}) {
+		fmt.Println("CreateTable Invitation")
+		db.Migrator().CreateTable(&Invitation{})
+	}
 	DB.MysqlDB.db = db
 	return
 }

pkg/common/db/mysql_model/im_mysql_model/invitation_model.go

@@ -0,0 +1,91 @@
+package im_mysql_model
+
+import (
+	"Open_IM/pkg/common/db"
+	"errors"
+	"math/rand"
+	"time"
+)
+
+/**
+ * 批量生成邀请码
+ */
+func BatchCreateInvitationCodes(CodeNums int, CodeLen int) error {
+	i := CodeNums
+	for {
+		if i == 0 {
+			break
+		}
+		invitation := new(db.Invitation)
+		invitation.CreateTime = time.Now()
+		invitation.InvitationCode = CreateRandomString(CodeLen)
+		invitation.LastTime = time.Now()
+		invitation.Status = 0
+		invitation.UserID = ""
+		result := db.DB.MysqlDB.DefaultGormDB().Table("invitations").Create(&invitation)
+		if result.Error != nil {
+			continue
+		}
+		if result.RowsAffected > 0 {
+			i = i - 1
+		}
+	}
+	return nil
+}
+
+/**
+ * 检查邀请码
+ */
+func CheckInvitationCode(code string) error {
+	var invitationCode db.Invitation
+	err := db.DB.MysqlDB.DefaultGormDB().Table("invitations").Where("invitation_code=?", code).Take(&invitationCode).Error
+	if err != nil {
+		return err
+	}
+	if invitationCode.InvitationCode != code {
+		return errors.New("邀请码不存在")
+	}
+	if invitationCode.Status != 0 {
+		return errors.New("邀请码已经被使用")
+	}
+	return nil
+}
+
+/**
+ * 尝试加锁模式解决邀请码抢占的问题
+ */
+func TryLockInvitationCode(Code string, UserId string) bool {
+	Data := make(map[string]interface{}, 0)
+	Data["user_id"] = UserId
+	Data["status"] = 1
+	Data["last_time"] = time.Now()
+	result := db.DB.MysqlDB.DefaultGormDB().Table("invitations").Where("invitation_code=? and user_id=? and status=?", Code, "", 0).Updates(Data)
+	if result.Error != nil {
+		return false
+	}
+	return result.RowsAffected > 0
+}
+
+/**
+ * 完成邀请码的状态
+ */
+func FinishInvitationCode(Code string, UserId string) bool {
+	Data := make(map[string]interface{}, 0)
+	Data["status"] = 2
+	result := db.DB.MysqlDB.DefaultGormDB().Table("invitations").Where("invitation_code=? and user_id=? and status=?", Code, UserId, 1).Updates(Data)
+	if result.Error != nil {
+		return false
+	}
+	return result.RowsAffected > 0
+}
+
+func CreateRandomString(strlen int) string {
+	str := "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+	bytes := []byte(str)
+	result := []byte{}
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	for i := 0; i < strlen; i++ {
+		result = append(result, bytes[r.Intn(len(bytes))])
+	}
+	return string(result)
+}

pkg/common/db/mysql_model/im_mysql_model/user_model.go

@@ -45,10 +45,20 @@ func UserRegister(user db.User) error {
 	user.LastLoginTime = time.Now()
 	user.LoginTimes = 0
 	user.LastLoginIp = user.CreateIp
+	if config.Config.Demo.NeedInvitationCode {
+		//判断一下验证码的使用情况
+		LockSucc := TryLockInvitationCode(user.InvitationCode, user.UserID)
+		if !LockSucc {
+			return constant.InvitationMsg
+		}
+	}
 	err := db.DB.MysqlDB.DefaultGormDB().Table("users").Create(&user).Error
 	if err != nil {
 		return err
 	}
+	if config.Config.Demo.NeedInvitationCode {
+		FinishInvitationCode(user.InvitationCode, user.UserID)
+	}
 	return nil
 }
 
@@ -125,11 +135,12 @@ func GetUsers(showNumber, pageNumber int32) ([]db.User, error) {
 
 func AddUser(userId, phoneNumber, name string) error {
 	user := db.User{
-		PhoneNumber: phoneNumber,
-		Birth:       time.Now(),
-		CreateTime:  time.Now(),
-		UserID:      userId,
-		Nickname:    name,
+		PhoneNumber:   phoneNumber,
+		Birth:         time.Now(),
+		CreateTime:    time.Now(),
+		UserID:        userId,
+		Nickname:      name,
+		LastLoginTime: time.Now(),
 	}
 	result := db.DB.MysqlDB.DefaultGormDB().Table("users").Create(&user)
 	return result.Error