Golang用户注册与权限控制实现方法

package main

import (
    "crypto/rand"
    "database/sql"
    "encoding/hex"
    "golang.org/x/crypto/bcrypt"
)

type User struct {
    ID       string
    Username string
    Password string
    Role     string
}

func hashPassword(password string) (string, error) {
    bytes, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
    return string(bytes), err
}

func generateToken() string {
    b := make([]byte, 16)
    rand.Read(b)
    return hex.EncodeToString(b)
}

func Register(db *sql.DB, username, password string) error {
    hashed, err := hashPassword(password)
    if err != nil {
        return err
    }

    token := generateToken()

    _, err = db.Exec(
        "INSERT INTO users (id, username, password, role, token) VALUES (?, ?, ?, ?, ?)",
        token, username, hashed, "user", token,
    )
    return err
}

func Login(db *sql.DB, username, password string) (string, error) {
    var user User
    err := db.QueryRow("SELECT id, password FROM users WHERE username = ?", username).
        Scan(&user.ID, &user.Password)
    if err != nil {
        return "", err
    }

    if err := bcrypt.CompareHashAndPassword([]byte(user.Password), []byte(password)); err != nil {
        return "", err
    }

    newToken := generateToken()
    _, err = db.Exec("UPDATE users SET token = ? WHERE username = ?", newToken, username)
    if err != nil {
        return "", err
    }

    return newToken, nil
}

func AuthMiddleware(db *sql.DB, requiredRole string) func(http.Handler) http.Handler {
    return func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
            token := r.Header.Get("Authorization")
            if token == "" {
                http.Error(w, "missing token", http.StatusUnauthorized)
                return
            }

            var user User
            err := db.QueryRow("SELECT id, username, role FROM users WHERE token = ?", token).
                Scan(&user.ID, &user.Username, &user.Role)
            if err != nil {
                http.Error(w, "invalid token", http.StatusUnauthorized)
                return
            }

            if requiredRole != "" && user.Role != requiredRole {
                http.Error(w, "forbidden", http.StatusForbidden)
                return
            }

            ctx := context.WithValue(r.Context(), "user", user)
            next.ServeHTTP(w, r.WithContext(ctx))
        })
    }
}

http.Handle("/api/profile", AuthMiddleware(db, "")(http.HandlerFunc(profileHandler)))
http.Handle("/api/admin", AuthMiddleware(db, "admin")(http.HandlerFunc(adminHandler)))