github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/service/jwt_token.go (about)

     1  package service
     2  
     3  import (
     4  	"fmt"
     5  	"time"
     6  
     7  	"github.com/golang-jwt/jwt"
     8  
     9  	"github.com/pyroscope-io/pyroscope/pkg/model"
    10  )
    11  
    12  const (
    13  	jwtClaimUserName = "name"
    14  	jwtClaimRole     = "role"
    15  )
    16  
    17  type JWTTokenService struct {
    18  	signingKey []byte
    19  	tokenTTL   time.Duration
    20  }
    21  
    22  func NewJWTTokenService(signingKey []byte, tokenTTL time.Duration) JWTTokenService {
    23  	return JWTTokenService{
    24  		signingKey: signingKey,
    25  		tokenTTL:   tokenTTL,
    26  	}
    27  }
    28  
    29  func (svc JWTTokenService) GenerateUserJWTToken(name string, role model.Role) *jwt.Token {
    30  	var exp time.Time
    31  	if svc.tokenTTL > 0 {
    32  		exp = time.Now().Add(svc.tokenTTL)
    33  	}
    34  	return generateToken(exp, jwt.MapClaims{
    35  		jwtClaimUserName: name,
    36  		jwtClaimRole:     role.String(),
    37  	})
    38  }
    39  
    40  func generateToken(exp time.Time, claims jwt.MapClaims) *jwt.Token {
    41  	claims["iat"] = time.Now().Unix()
    42  	if !exp.IsZero() {
    43  		claims["exp"] = exp.Unix()
    44  	}
    45  	return jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
    46  }
    47  
    48  // UserFromJWTToken retrieves user info from the given JWT token.
    49  // 'name' claim must be present and valid, otherwise the function returns
    50  // false. The function does not validate the token.
    51  func (JWTTokenService) UserFromJWTToken(t *jwt.Token) (model.TokenUser, bool) {
    52  	var user model.TokenUser
    53  	m, ok := t.Claims.(jwt.MapClaims)
    54  	if !ok {
    55  		return user, false
    56  	}
    57  	if user.Name, ok = m[jwtClaimUserName].(string); !ok {
    58  		return user, false
    59  	}
    60  	// Parse role.
    61  	s, ok := m[jwtClaimRole].(string)
    62  	if !ok {
    63  		return user, false
    64  	}
    65  	var err error
    66  	user.Role, err = model.ParseRole(s)
    67  	return user, err == nil
    68  }
    69  
    70  // Parse parses the token and validates it using the signing key.
    71  func (svc JWTTokenService) Parse(t string) (*jwt.Token, error) {
    72  	return jwt.Parse(t, func(token *jwt.Token) (interface{}, error) {
    73  		if token.Method.Alg() != jwt.SigningMethodHS256.Alg() {
    74  			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
    75  		}
    76  		return svc.signingKey, nil
    77  	})
    78  }
    79  
    80  func (svc JWTTokenService) Sign(t *jwt.Token) (string, error) {
    81  	return t.SignedString(svc.signingKey)
    82  }