erp_server/lib/auth/auth_user.go

package auth

import (
	"errors"

	"gitee.com/codinl/erp_server/lib/status"
	"strings"

	"github.com/codinl/go-logger"
	"github.com/dgrijalva/jwt-go"
	"github.com/gin-gonic/gin"
)

var (
	TokenExpired error = errors.New("token is expired")
	TokenInvalid error = errors.New("couldn't handle this token")
	//TokenNotValidYet error = errors.New("Token not active yet")
	//TokenMalformed   error = errors.New("That's not even a token")
)

var jwtUserAccessKey string
var jwtUserRefreshKey string
var jwtUserAccessExpire int64
var jwtUserRefreshExpire int64

var jwtAdminAccessKey string

//var jwtAdminRefreshKey string
//var jwtAdminAccessExpire int64
//var jwtAdminRefreshExpire int64

func Init(userAccessKey string, userAccessExpire int64) {
	jwtUserAccessKey = userAccessKey
	jwtUserAccessExpire = userAccessExpire
}

func UserAccessAuth(c *gin.Context) {
	token := c.Request.Header.Get("Authorization")
	if s := strings.Split(token, " "); len(s) == 2 {
		token = s[1]
	}
	if token == "" || len(token) < 10 {
		// TODO
		//RespJson(c, status.Unauthorized, nil)
		return
	}

	j := NewJWT(jwtUserAccessKey)
	claims, err := j.ParseToken(token)
	if err != nil {
		logger.Error(err)
		if err == TokenExpired {
			RespJson(c, status.AccessTokenExpired, nil)
		}
		//else { 		// TODO
		//	RespJson(c, status.Unauthorized, nil)
		//}
		return
	}

	c.Set("claims", claims)
}

func AdminAccessAuth(c *gin.Context) {
	token := c.Request.Header.Get("Authorization")
	if s := strings.Split(token, " "); len(s) == 2 {
		token = s[1]
	}
	if token == "" || len(token) < 10 {
		RespJson(c, status.Unauthorized, nil)
		return
	}

	j := NewJWT(jwtAdminAccessKey)
	claims, err := j.ParseToken(token)
	if err != nil {
		RespJson(c, status.Unauthorized, nil)
		return
	}

	c.Set("claims", claims)
}

func RespJson(c *gin.Context, code int, data interface{}) {
	result := struct {
		Code int         `json:"code"`
		Msg  string      `json:"msg"`
		Desc string      `json:"desc"`
		Data interface{} `json:"data"`
	}{
		Code: code,
		Desc: status.StatusDesc(code),
		Data: data,
	}
	c.JSON(status.OK, result)
	c.Abort()
}

type JWT struct {
	SigningKey []byte
}

type UserClaims struct {
	Uid uint32 `json:"uid"`
	jwt.StandardClaims
}

func NewJWT(key string) *JWT {
	return &JWT{
		[]byte(key),
	}
}

//func (j *JWT) CreateToken(user *UserClaims) (string, error) {
//	user.ExpiresAt = jwt.TimeFunc().Unix() + int64(jwtUserAccessExpire)
//	token := jwt.NewWithClaims(jwt.SigningMethodHS256, user)
//	return token.SignedString(j.SigningKey)
//}

func (j *JWT) CreateToken(uid uint32, key string, expire int64) (string, error) {
	user := UserClaims{
		Uid: uid,
	}
	user.ExpiresAt = jwt.TimeFunc().Unix() + expire
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, user)
	return token.SignedString(key)
}

func (j *JWT) ParseToken(tokenString string) (*UserClaims, error) {
	token, err := jwt.ParseWithClaims(tokenString, &UserClaims{}, func(token *jwt.Token) (interface{}, error) {
		return j.SigningKey, nil
	})

	if err != nil {
		e := err.(*jwt.ValidationError)
		if e.Errors&jwt.ValidationErrorExpired != 0 {
			return nil, TokenExpired
		}
		return nil, err
	}

	if claims, ok := token.Claims.(*UserClaims); ok && token.Valid {
		return claims, nil
	} else {
		logger.Error(err)
	}

	return nil, TokenInvalid
}

func GetCurrentUser(c *gin.Context) *UserClaims {
	claims, ok := c.Get("claims")
	if !ok {
		return nil
	}
	user, ok := claims.(*UserClaims)
	if ok {
		return user
	}
	return nil
}

func CreateToken(uid uint32, key string, expire int64) (string, error) {
	user := UserClaims{
		Uid: uid,
	}
	user.ExpiresAt = jwt.TimeFunc().Unix() + expire
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, user)
	return token.SignedString([]byte(key))
}

func CreateAccessToken(uid uint32) (string, error) {
	return CreateToken(uid, jwtUserAccessKey, jwtUserAccessExpire)
}

func CreateRefreshToken(uid uint32) (string, error) {
	return CreateToken(uid, jwtUserRefreshKey, jwtUserRefreshExpire)
}