356 lines
10 KiB
Go
356 lines
10 KiB
Go
package repository
|
||
|
||
import (
|
||
"database/sql"
|
||
"errors"
|
||
"fmt"
|
||
"strings"
|
||
|
||
"github.com/wmms/wmms/internal/model"
|
||
"github.com/wmms/wmms/pkg/database"
|
||
"golang.org/x/crypto/bcrypt"
|
||
)
|
||
|
||
var (
|
||
ErrInvalidRoles = errors.New("one or more roles do not exist or are disabled")
|
||
ErrLastSuperAdmin = errors.New("the last super administrator cannot be removed")
|
||
)
|
||
|
||
type UserRepository struct{}
|
||
|
||
func NewUserRepository() *UserRepository { return &UserRepository{} }
|
||
|
||
func (r *UserRepository) GetByUsername(username string) (*model.User, error) {
|
||
user := &model.User{}
|
||
err := database.DB.QueryRow(`SELECT id, username, password, description, created_at FROM users WHERE username = ?`, username).Scan(
|
||
&user.ID, &user.Username, &user.Password, &user.Description, &user.CreatedAt,
|
||
)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if err := loadRolesForUsers(database.DB, []*model.User{user}); err != nil {
|
||
return nil, err
|
||
}
|
||
return user, nil
|
||
}
|
||
|
||
func (r *UserRepository) GetByID(id int64) (*model.User, error) {
|
||
user := &model.User{}
|
||
err := database.DB.QueryRow(`SELECT id, username, password, description, created_at FROM users WHERE id = ?`, id).Scan(
|
||
&user.ID, &user.Username, &user.Password, &user.Description, &user.CreatedAt,
|
||
)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if err := loadRolesForUsers(database.DB, []*model.User{user}); err != nil {
|
||
return nil, err
|
||
}
|
||
return user, nil
|
||
}
|
||
|
||
func (r *UserRepository) Create(user *model.User, roleIDs []int64) error {
|
||
tx, err := database.DB.Begin()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer func() { _ = tx.Rollback() }()
|
||
|
||
result, err := tx.Exec(`INSERT INTO users (username, password, description) VALUES (?, ?, ?)`, user.Username, user.Password, user.Description)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
user.ID, err = result.LastInsertId()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if err = assignRolesTx(tx, user.ID, roleIDs); err != nil {
|
||
return err
|
||
}
|
||
return tx.Commit()
|
||
}
|
||
|
||
func (r *UserRepository) List(search string, limit, offset int) ([]model.User, int64, error) {
|
||
countQuery := `SELECT COUNT(*) FROM users`
|
||
var countArgs []interface{}
|
||
if search != "" {
|
||
countQuery += ` WHERE username LIKE ?`
|
||
countArgs = append(countArgs, "%"+search+"%")
|
||
}
|
||
var total int64
|
||
if err := database.DB.QueryRow(countQuery, countArgs...).Scan(&total); err != nil {
|
||
return nil, 0, err
|
||
}
|
||
|
||
query := `SELECT id, username, password, description, created_at FROM users`
|
||
var args []interface{}
|
||
if search != "" {
|
||
query += ` WHERE username LIKE ?`
|
||
args = append(args, "%"+search+"%")
|
||
}
|
||
query += ` ORDER BY id LIMIT ? OFFSET ?`
|
||
args = append(args, limit, offset)
|
||
rows, err := database.DB.Query(query, args...)
|
||
if err != nil {
|
||
return nil, 0, err
|
||
}
|
||
defer rows.Close()
|
||
|
||
users := make([]model.User, 0)
|
||
for rows.Next() {
|
||
var user model.User
|
||
if err := rows.Scan(&user.ID, &user.Username, &user.Password, &user.Description, &user.CreatedAt); err != nil {
|
||
return nil, 0, err
|
||
}
|
||
users = append(users, user)
|
||
}
|
||
if err := rows.Err(); err != nil {
|
||
return nil, 0, err
|
||
}
|
||
userPointers := make([]*model.User, 0, len(users))
|
||
for index := range users {
|
||
userPointers = append(userPointers, &users[index])
|
||
}
|
||
if err := loadRolesForUsers(database.DB, userPointers); err != nil {
|
||
return nil, 0, err
|
||
}
|
||
return users, total, nil
|
||
}
|
||
|
||
func (r *UserRepository) Update(user *model.User, roleIDs []int64) error {
|
||
return r.update(user, roleIDs, false)
|
||
}
|
||
|
||
func (r *UserRepository) UpdateWithPassword(user *model.User, roleIDs []int64) error {
|
||
return r.update(user, roleIDs, true)
|
||
}
|
||
|
||
func (r *UserRepository) update(user *model.User, roleIDs []int64, updatePassword bool) error {
|
||
tx, err := database.DB.Begin()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer func() { _ = tx.Rollback() }()
|
||
|
||
roleIDs = uniqueRoleIDs(roleIDs)
|
||
if err := validateRolesTx(tx, user.ID, roleIDs); err != nil {
|
||
return err
|
||
}
|
||
if err := protectLastSuperAdminTx(tx, user.ID, roleIDs); err != nil {
|
||
return err
|
||
}
|
||
if updatePassword {
|
||
_, err = tx.Exec(`UPDATE users SET password = ?, description = ? WHERE id = ?`, user.Password, user.Description, user.ID)
|
||
} else {
|
||
_, err = tx.Exec(`UPDATE users SET description = ? WHERE id = ?`, user.Description, user.ID)
|
||
}
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if _, err = tx.Exec(`DELETE FROM user_roles WHERE user_id = ?`, user.ID); err != nil {
|
||
return err
|
||
}
|
||
if err = insertUserRolesTx(tx, user.ID, roleIDs); err != nil {
|
||
return err
|
||
}
|
||
return tx.Commit()
|
||
}
|
||
|
||
func (r *UserRepository) Delete(id int64) error {
|
||
tx, err := database.DB.Begin()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer func() { _ = tx.Rollback() }()
|
||
if err := protectLastSuperAdminTx(tx, id, nil); err != nil {
|
||
return err
|
||
}
|
||
if _, err := tx.Exec(`DELETE FROM users WHERE id = ?`, id); err != nil {
|
||
return err
|
||
}
|
||
return tx.Commit()
|
||
}
|
||
|
||
func loadRolesForUsers(queryer interface {
|
||
Query(query string, args ...interface{}) (*sql.Rows, error)
|
||
}, users []*model.User) error {
|
||
if len(users) == 0 {
|
||
return nil
|
||
}
|
||
indexes := make(map[int64]*model.User, len(users))
|
||
args := make([]interface{}, 0, len(users))
|
||
for _, user := range users {
|
||
user.Roles = []model.Role{}
|
||
indexes[user.ID] = user
|
||
args = append(args, user.ID)
|
||
}
|
||
placeholders := strings.TrimSuffix(strings.Repeat("?, ", len(args)), ", ")
|
||
query := fmt.Sprintf(`SELECT ur.user_id, r.id, r.code, r.name, r.enabled, r.created_at, r.updated_at
|
||
FROM user_roles ur JOIN roles r ON r.id = ur.role_id
|
||
WHERE ur.user_id IN (%s) ORDER BY ur.user_id, r.id`, placeholders)
|
||
rows, err := queryer.Query(query, args...)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer rows.Close()
|
||
for rows.Next() {
|
||
var userID int64
|
||
var role model.Role
|
||
if err := rows.Scan(&userID, &role.ID, &role.Code, &role.Name, &role.Enabled, &role.CreatedAt, &role.UpdatedAt); err != nil {
|
||
return err
|
||
}
|
||
role.Permissions = []model.Permission{}
|
||
if user := indexes[userID]; user != nil {
|
||
user.Roles = append(user.Roles, role)
|
||
}
|
||
}
|
||
return rows.Err()
|
||
}
|
||
|
||
func assignRolesTx(tx *sql.Tx, userID int64, roleIDs []int64) error {
|
||
roleIDs = uniqueRoleIDs(roleIDs)
|
||
if err := validateRolesTx(tx, userID, roleIDs); err != nil {
|
||
return err
|
||
}
|
||
return insertUserRolesTx(tx, userID, roleIDs)
|
||
}
|
||
|
||
// validateRolesTx 校验 roleIDs 是否都是可分配给该用户的角色。
|
||
// 可分配 = 角色处于启用状态，或它本来就是该用户当前持有的角色。
|
||
// - 新建用户时 user_roles 里还没有该用户的记录，EXISTS 恒为假，等价于原来的
|
||
// "只允许启用中的角色"。
|
||
// - 编辑用户时前端会把现有角色原样回传，其中可能含已被停用的。停用只是不再授予
|
||
// 权限（GenerateToken 会过滤掉停用角色），没必要把已有分配一并抹掉；若在这里
|
||
// 拒绝，管理员就再也保存不了这个用户：下拉里的停用角色是禁用项，已选中的 tag
|
||
// 也没有关闭按钮。
|
||
//
|
||
// 必须在 DELETE FROM user_roles 之前调用，否则 EXISTS 取不到旧分配。
|
||
func validateRolesTx(tx *sql.Tx, userID int64, roleIDs []int64) error {
|
||
if len(roleIDs) == 0 {
|
||
return ErrInvalidRoles
|
||
}
|
||
placeholders := strings.TrimSuffix(strings.Repeat("?, ", len(roleIDs)), ", ")
|
||
args := make([]interface{}, 0, len(roleIDs)+1)
|
||
for _, id := range roleIDs {
|
||
args = append(args, id)
|
||
}
|
||
args = append(args, userID)
|
||
var count int
|
||
query := fmt.Sprintf(`SELECT COUNT(*) FROM roles r
|
||
WHERE r.id IN (%s)
|
||
AND (r.enabled = TRUE OR EXISTS (
|
||
SELECT 1 FROM user_roles ur WHERE ur.role_id = r.id AND ur.user_id = ?))`, placeholders)
|
||
if err := tx.QueryRow(query, args...).Scan(&count); err != nil {
|
||
return err
|
||
}
|
||
if count != len(roleIDs) {
|
||
return ErrInvalidRoles
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func insertUserRolesTx(tx *sql.Tx, userID int64, roleIDs []int64) error {
|
||
values := strings.TrimSuffix(strings.Repeat("(?, ?), ", len(roleIDs)), ", ")
|
||
args := make([]interface{}, 0, len(roleIDs)*2)
|
||
for _, roleID := range roleIDs {
|
||
args = append(args, userID, roleID)
|
||
}
|
||
_, err := tx.Exec(`INSERT INTO user_roles (user_id, role_id) VALUES `+values, args...)
|
||
return err
|
||
}
|
||
|
||
func protectLastSuperAdminTx(tx *sql.Tx, userID int64, replacementRoleIDs []int64) error {
|
||
var superAdminRoleID int64
|
||
if err := tx.QueryRow(`SELECT id FROM roles WHERE code = 'super_admin' FOR UPDATE`).Scan(&superAdminRoleID); err != nil {
|
||
return err
|
||
}
|
||
for _, roleID := range replacementRoleIDs {
|
||
if roleID == superAdminRoleID {
|
||
return nil
|
||
}
|
||
}
|
||
var currentMembership int
|
||
if err := tx.QueryRow(`SELECT COUNT(*) FROM user_roles WHERE user_id = ? AND role_id = ?`, userID, superAdminRoleID).Scan(¤tMembership); err != nil {
|
||
return err
|
||
}
|
||
if currentMembership == 0 {
|
||
return nil
|
||
}
|
||
var superAdminCount int
|
||
if err := tx.QueryRow(`SELECT COUNT(*) FROM user_roles WHERE role_id = ?`, superAdminRoleID).Scan(&superAdminCount); err != nil {
|
||
return err
|
||
}
|
||
if superAdminCount <= 1 {
|
||
return ErrLastSuperAdmin
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func uniqueRoleIDs(roleIDs []int64) []int64 {
|
||
seen := make(map[int64]struct{}, len(roleIDs))
|
||
result := make([]int64, 0, len(roleIDs))
|
||
for _, id := range roleIDs {
|
||
if id <= 0 {
|
||
continue
|
||
}
|
||
if _, exists := seen[id]; exists {
|
||
continue
|
||
}
|
||
seen[id] = struct{}{}
|
||
result = append(result, id)
|
||
}
|
||
return result
|
||
}
|
||
|
||
func (r *UserRepository) UserExists(username string) (bool, error) {
|
||
var count int
|
||
err := database.DB.QueryRow(`SELECT COUNT(*) FROM users WHERE username = ?`, username).Scan(&count)
|
||
return count > 0, err
|
||
}
|
||
|
||
func InitUserTable() error {
|
||
_, err := database.DB.Exec(`CREATE TABLE IF NOT EXISTS users (
|
||
id BIGINT AUTO_INCREMENT PRIMARY KEY,
|
||
username VARCHAR(50) NOT NULL UNIQUE,
|
||
password VARCHAR(255) NOT NULL,
|
||
description VARCHAR(500) DEFAULT '',
|
||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
||
)`)
|
||
return err
|
||
}
|
||
|
||
func CreateDefaultUsers() error {
|
||
repo := NewUserRepository()
|
||
exists, err := repo.UserExists("admin")
|
||
if err != nil || exists {
|
||
return err
|
||
}
|
||
adminRoleID, err := roleIDByCode("super_admin")
|
||
if err != nil {
|
||
return err
|
||
}
|
||
warehouseRoleID, err := roleIDByCode("warehouse_admin")
|
||
if err != nil {
|
||
return err
|
||
}
|
||
adminHash, err := bcrypt.GenerateFromPassword([]byte("admin123"), bcrypt.DefaultCost)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
admin := &model.User{Username: "admin", Password: string(adminHash), Description: "系统管理员"}
|
||
if err := repo.Create(admin, []int64{adminRoleID}); err != nil {
|
||
return err
|
||
}
|
||
userHash, err := bcrypt.GenerateFromPassword([]byte("user123"), bcrypt.DefaultCost)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
user := &model.User{Username: "user", Password: string(userHash), Description: "普通用户"}
|
||
return repo.Create(user, []int64{warehouseRoleID})
|
||
}
|
||
|
||
func roleIDByCode(code string) (int64, error) {
|
||
var id int64
|
||
err := database.DB.QueryRow(`SELECT id FROM roles WHERE code = ? AND enabled = TRUE`, code).Scan(&id)
|
||
return id, err
|
||
}
|