PrimeChatServer/user/method.go

package user

import (
	"errors"
	"fmt"
	"time"

	. "prime/basic"

	"github.com/golang-jwt/jwt/v4"
	"gorm.io/gorm"
)

type MyClaims struct {
	Username string `json:"username"`
	UID      uint   `json:"uid"`
	jwt.StandardClaims
}

const TokenExpireDuration = time.Hour * 2

var MySecret = []byte("冬天冬天悄悄过去")

func GenToken(username string, password string) (string, error) {
	db := DB()
	var user User
	result := db.Take(&user, "username = ? AND password = ?", username, password)
	fmt.Printf("%+v\n", user)
	if result.RowsAffected == 1 {
		c := MyClaims{username, user.ID, jwt.StandardClaims{
			ExpiresAt: time.Now().Add(TokenExpireDuration).Unix(),
			Issuer:    "my-project",
		}}
		token := jwt.NewWithClaims(jwt.SigningMethodHS256, c)
		return token.SignedString(MySecret)
	}
	return "", gorm.ErrRecordNotFound
}

func ParseToken(tokenString string) (*MyClaims, error) {
	// 解析token
	token, err := jwt.ParseWithClaims(tokenString, &MyClaims{}, func(token *jwt.Token) (i interface{}, err error) {
		return MySecret, nil
	})
	if err != nil {
		return nil, err
	}
	if claims, ok := token.Claims.(*MyClaims); ok && token.Valid { // 校验token
		return claims, nil
	}
	return nil, errors.New("invalid token")
}

func SetAvatar(uid uint, filename string) error {
	db := DB()
	var user User
	user.ID = uid
	var profile Profile
	err := db.Model(&user).Association("Profile").Find(&profile)
	if err != nil {
		return err
	}
	profile.Avatar = filename
	fmt.Println(profile)
	res := db.Where("user_id = ?", uid).Save(&profile)
	// return nil
	return res.Error
}

func GetAvatar(uid uint) (string, error) {
	db := DB()
	var user User
	user.ID = uid
	var profile Profile
	err := db.Model(&user).Association("Profile").Find(&profile)
	if err != nil {
		return "", err
	}
	return profile.Avatar, nil
}

func Register(username string, password string) (user *User, err error) {
	db := DB()
	profile := Profile{Name: username, Avatar: "avatar.jpg"}
	user = &User{Username: username, Password: password, Profile: profile}
	result := db.Create(user)
	return user, result.Error
}

func GetProfile(uid uint) (*Profile, error) {
	db := DB()
	var profile Profile
	res := db.Find(&profile, "user_id = ?", uid)
	if res.RowsAffected == 1 {
		return &profile, nil
	}
	return nil, res.Error
}

func GetUID(username string) (uint, error) {
	db := DB()
	var user User
	res := db.Find(&user, "username = ?", username)
	if res.RowsAffected == 1 {
		return user.ID, nil
	}
	if res.Error == gorm.ErrRecordNotFound {
		println("user does not exist")
		return 0, nil
	}
	return 0, res.Error
}