bj_power/bj_power_wms/internal/handler/auth.go

package handler

import (
	"net/http"
	"strings"
	"time"

	"bj_power_wms/ent/user"
	"bj_power_wms/internal/svc"

	"golang.org/x/crypto/bcrypt"
)

func healthHandler(ctx *svc.ServiceContext) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		_, err := ctx.EntClient.User.Query().First(ctx0())
		if err == nil || strings.Contains(err.Error(), "no rows") {
			ok(w, map[string]any{
				"status":  "UP",
				"time":    time.Now().Format("2006-01-02 15:04:05"),
				"service": "bj_power_wms",
			})
			return
		}
		fail(w, http.StatusServiceUnavailable, "database unavailable")
	}
}

func loginHandler(ctx *svc.ServiceContext) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var req struct {
			Username string `json:"username"`
			Password string `json:"password"`
		}
		if err := parseJSON(r, &req); err != nil {
			fail(w, http.StatusBadRequest, "参数错误")
			return
		}
		if req.Username == "" || req.Password == "" {
			fail(w, http.StatusBadRequest, "用户名和密码必填")
			return
		}

		u, err := ctx.EntClient.User.Query().
			Where(user.UsernameEQ(req.Username)).
			Only(ctx0())
		if err != nil {
			fail(w, http.StatusUnauthorized, "用户名或密码错误")
			return
		}
		if !u.IsActive {
			fail(w, http.StatusForbidden, "账号已停用")
			return
		}
		if bcrypt.CompareHashAndPassword([]byte(u.Password), []byte(req.Password)) != nil {
			fail(w, http.StatusUnauthorized, "用户名或密码错误")
			return
		}

		ctx.EntClient.User.UpdateOneID(u.ID).SetLastLoginAt(time.Now().Unix()).ExecX(ctx0())

		token, expireAt, err := SignToken(ctx, u.ID, u.Username, u.RealName, u.Role)
		if err != nil {
			fail(w, http.StatusInternalServerError, "签发 token 失败")
			return
		}
		ok(w, map[string]any{
			"token":    token,
			"expireAt": expireAt,
			"user":     map[string]any{"id": u.ID, "username": u.Username, "realName": u.RealName, "role": u.Role},
		})
	}
}

func registerHandler(ctx *svc.ServiceContext) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var req struct {
			Username string `json:"username"`
			Password string `json:"password"`
			RealName string `json:"realName"`
			Role     string `json:"role"`
		}
		if err := parseJSON(r, &req); err != nil {
			fail(w, http.StatusBadRequest, "参数错误")
			return
		}
		if req.Username == "" || req.Password == "" {
			fail(w, http.StatusBadRequest, "用户名和密码必填")
			return
		}
		if len(req.Password) < 6 {
			fail(w, http.StatusBadRequest, "密码至少 6 位")
			return
		}
		exists, _ := ctx.EntClient.User.Query().Where(user.UsernameEQ(req.Username)).Exist(ctx0())
		if exists {
			fail(w, http.StatusConflict, "用户名已存在")
			return
		}
		hash, err := bcrypt.GenerateFromPassword([]byte(req.Password), bcrypt.DefaultCost)
		if err != nil {
			fail(w, http.StatusInternalServerError, err.Error())
			return
		}
		u, err := ctx.EntClient.User.Create().
			SetUsername(req.Username).
			SetPassword(string(hash)).
			SetNillableRealName(strPtr(req.RealName)).
			SetRole(roleOrDefault(req.Role)).
			Save(ctx0())
		if err != nil {
			fail(w, http.StatusInternalServerError, err.Error())
			return
		}
		ok(w, map[string]any{"id": u.ID, "username": u.Username})
	}
}

func roleOrDefault(r string) string {
	switch r {
	case "admin", "operator", "inspector":
		return r
	default:
		return "operator"
	}
}