342 lines
9.6 KiB
Go
342 lines
9.6 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, 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, roleIDs); err != nil {
|
|
return err
|
|
}
|
|
return insertUserRolesTx(tx, userID, roleIDs)
|
|
}
|
|
|
|
func validateRolesTx(tx *sql.Tx, roleIDs []int64) error {
|
|
if len(roleIDs) == 0 {
|
|
return ErrInvalidRoles
|
|
}
|
|
placeholders := strings.TrimSuffix(strings.Repeat("?, ", len(roleIDs)), ", ")
|
|
args := make([]interface{}, len(roleIDs))
|
|
for index, id := range roleIDs {
|
|
args[index] = id
|
|
}
|
|
var count int
|
|
query := fmt.Sprintf(`SELECT COUNT(*) FROM roles WHERE enabled = TRUE AND id IN (%s)`, 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
|
|
}
|