licensing-cotton/internal/handlers/auth.go
2025-01-16 15:41:42 +08:00

118 lines
3.2 KiB
Go

package handlers
import (
"database/sql"
"encoding/json"
"licensing-cotton/internal/security"
"net/http"
"strings"
"golang.org/x/crypto/bcrypt"
"licensing-cotton/internal/database"
"licensing-cotton/internal/models"
)
func HandleLogin(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
http.Error(w, "只允许POST", http.StatusMethodNotAllowed)
return
}
var req struct {
Username string `json:"username"`
Password string `json:"password"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
http.Error(w, "请求解析失败", http.StatusBadRequest)
return
}
// 从数据库查询用户
var user models.User
err := database.DB.QueryRow(`SELECT id, username, password_hash, role
FROM users WHERE username = ?`, req.Username).
Scan(&user.ID, &user.Username, &user.PasswordHash, &user.Role)
if err == sql.ErrNoRows {
http.Error(w, "用户名或密码错误", http.StatusUnauthorized)
return
} else if err != nil {
http.Error(w, "数据库查询错误", http.StatusInternalServerError)
return
}
// 验证密码
if err := bcrypt.CompareHashAndPassword([]byte(user.PasswordHash), []byte(req.Password)); err != nil {
http.Error(w, "用户名或密码错误", http.StatusUnauthorized)
return
}
// 生成 session
token, _ := security.GenerateToken(user.Username) // session.go 或本文件里实现
setSession(token, user.Username)
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]string{
"message": "登录成功",
"token": token,
"role": user.Role,
})
}
// 注册(仅管理员能创建/更新用户)
func HandleRegisterUser(w http.ResponseWriter, r *http.Request) {
if !isAdminRequest(w, r) {
return
}
if r.Method != http.MethodPost {
http.Error(w, "只允许POST", http.StatusMethodNotAllowed)
return
}
var req struct {
Username string `json:"username"`
Password string `json:"password"`
Role string `json:"role"` // "admin" or "user"
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
http.Error(w, "解析请求失败", http.StatusBadRequest)
return
}
passHash, _ := bcrypt.GenerateFromPassword([]byte(req.Password), bcrypt.DefaultCost)
_, err := database.DB.Exec(`INSERT INTO users(username, password_hash, role) VALUES (?,?,?)`,
req.Username, passHash, req.Role)
if err != nil {
if strings.Contains(err.Error(), "UNIQUE constraint failed") {
http.Error(w, "用户名已存在", http.StatusConflict)
} else {
http.Error(w, "插入数据库失败: "+err.Error(), http.StatusInternalServerError)
}
return
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]string{
"message": "用户创建成功",
})
}
// 初始化默认管理员
func EnsureDefaultAdmin(username, password string) error {
var count int
err := database.DB.QueryRow(`SELECT COUNT(*) FROM users WHERE username = ?`, username).Scan(&count)
if err != nil {
return err
}
if count == 0 {
passHash, _ := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
_, err := database.DB.Exec(`INSERT INTO users(username, password_hash, role) VALUES (?,?,?)`,
username, passHash, "admin")
if err != nil {
return err
}
}
return nil
}