#4#8

Author: kuangpenghao

.github/workflows/docker-image.yml

@@ -1,7 +1,7 @@
 name: Docker Image CI
 on:
   push:
-    branches: [ "master" ]
+    branches: [ "**" ]
 
 jobs:
   build:

internal/controllers/users/register.go

@@ -5,16 +5,18 @@ import (
 	"coursebench-backend/pkg/errors"
 	"coursebench-backend/pkg/models"
 	"coursebench-backend/pkg/queries"
+
 	"github.com/gofiber/fiber/v2"
 )
 
 type RegisterRequest struct {
-	Email    string           `json:"email"`
-	Password string           `json:"password"`
-	Year     int              `json:"year"`
-	Grade    models.GradeType `json:"grade"`
-	Captcha  string           `json:"captcha"`
-	Nickname string           `json:"nickname"`
+	Email          string           `json:"email"`
+	Password       string           `json:"password"`
+	Year           int              `json:"year"`
+	Grade          models.GradeType `json:"grade"`
+	Captcha        string           `json:"captcha"`
+	Nickname       string           `json:"nickname"`
+	InvitationCode string           `json:"invitation_code"`
 }
 
 func Register(c *fiber.Ctx) (err error) {
@@ -38,7 +40,7 @@ func Register(c *fiber.Ctx) (err error) {
 		Avatar:      "",
 		IsAnonymous: false,
 	}
-	if err = queries.Register(nil, &user); err != nil {
+	if err = queries.Register(nil, &user, userReq.InvitationCode); err != nil {
 		return
 	}
 	if config.GlobalConf.DisableMail {

pkg/errors/description.go

@@ -33,6 +33,7 @@ var (
 	CaptchaMismatch       = createDescription("CaptchaMismatch", "验证码错误", SILENT, 400)
 	NoCaptchaToken        = createDescription("NoCaptchaToken", "未请求过验证码Token，请检查您的 Cookie 设置", SILENT, 400)
 	CaptchaExpired        = createDescription("CaptchaExpired", "验证码已过期", SILENT, 400)
+	InvitationCodeInvalid = createDescription("InvitationCodeInvalid", "邀请码无效", SILENT, 400)
 
 	TeacherNotExists = createDescription("TeacherNotExists", "未找到教师", SILENT, 400)
 

pkg/models/user.go

@@ -2,6 +2,7 @@ package models
 
 import (
 	"coursebench-backend/pkg/modelRegister"
+
 	"gorm.io/gorm"
 )
 
@@ -47,4 +48,5 @@ type ProfileResponse struct {
 	IsAnonymous      bool      `json:"is_anonymous"`
 	IsAdmin          bool      `json:"is_admin"`
 	IsCommunityAdmin bool      `json:"is_community_admin"`
+	InvitationCode   string    `json:"invitation_code"`
 }

pkg/queries/user.go

@@ -7,12 +7,14 @@ import (
 	"coursebench-backend/pkg/mail"
 	"coursebench-backend/pkg/models"
 	"fmt"
-	"github.com/badoux/checkmail"
-	"golang.org/x/crypto/bcrypt"
-	"gorm.io/gorm"
+	"math/rand"
 	"strings"
 	"time"
 	"unicode"
+
+	"github.com/badoux/checkmail"
+	"golang.org/x/crypto/bcrypt"
+	"gorm.io/gorm"
 )
 
 func ResetPassword(db *gorm.DB, email string) error {
@@ -71,7 +73,7 @@ func ResetPasswordActive(db *gorm.DB, id uint, code string, password string) (er
 	return nil
 }
 
-func Register(db *gorm.DB, u *models.User) error {
+func Register(db *gorm.DB, u *models.User, invitation_code string) error {
 	if db == nil {
 		db = database.GetDB()
 	}
@@ -95,6 +97,22 @@ func Register(db *gorm.DB, u *models.User) error {
 	if !CheckRealName(u.RealName) {
 		return errors.New(errors.InvalidArgument)
 	}
+	if !CheckInvitationCode(invitation_code) {
+		return errors.New(errors.InvalidArgument)
+	}
+
+	// check if the invitation code is valid
+	if invitation_code != "" {
+		taken, err := isInvitationCodeTaken(db, invitation_code)
+		if err != nil {
+			return err
+		}
+		if !taken {
+			return errors.New(errors.InvitationCodeInvalid)
+		}
+
+		// TODO: Inform the inviter
+	}
 
 	// 检查邮箱是否已存在
 	user := &models.User{}
@@ -117,6 +135,12 @@ func Register(db *gorm.DB, u *models.User) error {
 	u.IsActive = false
 	u.IsAdmin = false
 
+	code, err := createInvitationCode(db)
+	if err != nil {
+		return err
+	}
+	u.InvitationCode = code
+
 	if err = db.Create(u).Error; err != nil {
 		return errors.Wrap(err, errors.DatabaseError)
 	}
@@ -233,6 +257,19 @@ func Login(db *gorm.DB, email, password string) (*models.User, error) {
 		return nil, errors.New(errors.UserNotActive)
 	}
 
+	if user.InvitationCode == "" {
+		code, err := createInvitationCode(db)
+		if err != nil {
+			return nil, err
+		}
+
+		user.InvitationCode = code
+		err = db.Select("invitation_code").Save(user).Error
+		if err != nil {
+			return nil, errors.Wrap(err, errors.DatabaseError)
+		}
+	}
+
 	return user, nil
 }
 
@@ -332,7 +369,11 @@ func GetProfile(db *gorm.DB, id uint, uid uint) (models.ProfileResponse, error)
 	if user.IsAnonymous && id != uid {
 		return models.ProfileResponse{ID: id, NickName: user.NickName, Avatar: avatar, IsAnonymous: user.IsAnonymous, IsAdmin: user.IsAdmin, IsCommunityAdmin: user.IsCommunityAdmin}, nil
 	} else {
-		return models.ProfileResponse{ID: id, Email: user.Email, Year: user.Year, Grade: user.Grade, NickName: user.NickName, RealName: user.RealName, IsAnonymous: user.IsAnonymous, Avatar: avatar, IsAdmin: user.IsAdmin, IsCommunityAdmin: user.IsCommunityAdmin}, nil
+		r := models.ProfileResponse{ID: id, Email: user.Email, Year: user.Year, Grade: user.Grade, NickName: user.NickName, RealName: user.RealName, IsAnonymous: user.IsAnonymous, Avatar: avatar, IsAdmin: user.IsAdmin, IsCommunityAdmin: user.IsCommunityAdmin}
+		if id == uid {
+			r.InvitationCode = user.InvitationCode
+		}
+		return r, nil
 	}
 }
 
@@ -397,3 +438,48 @@ func CheckRealName(realname string) bool {
 	}
 	return true
 }
+
+func CheckInvitationCode(code string) bool {
+	if len(code) == 0 {
+		return true
+	}
+	if len(code) != 5 {
+		return false
+	}
+	for _, c := range code {
+		if (c < '0' || c > '9') && (c < 'a' || c > 'z') && (c < 'A' || c > 'Z') {
+			return false
+		}
+	}
+	return true
+}
+
+func createInvitationCode(db *gorm.DB) (string, error) {
+	// try a few times before giving up
+	for i := 0; i < 5; i++ {
+		codeRunes := make([]rune, 0, 5)
+		for i := 0; i < 5; i++ {
+			codeRunes = append(codeRunes, []rune("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789")[rand.Intn(62)])
+		}
+		code := string(codeRunes)
+
+		taken, err := isInvitationCodeTaken(db, code)
+		if err != nil {
+			return "", err
+		}
+		if !taken {
+			return code, nil
+		}
+	}
+
+	return "", errors.New(errors.InternalServerError)
+}
+
+func isInvitationCodeTaken(db *gorm.DB, code string) (bool, error) {
+	user := &models.User{}
+	result := db.Where("invitation_code = ?", code).Take(user)
+	if err := result.Error; err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
+		return false, errors.Wrap(err, errors.DatabaseError)
+	}
+	return result.RowsAffected != 0, nil
+}