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

     1  package storage
     2  
     3  import (
     4  	"crypto/rand"
     5  	"math/big"
     6  
     7  	"github.com/dgraph-io/badger/v2"
     8  )
     9  
    10  const (
    11  	jwtLenght = 32
    12  	jwtSecret = "jwtSecret"
    13  )
    14  
    15  func (s *Storage) JWT() (string, error) {
    16  	var secret []byte
    17  	err := s.main.View(func(txn *badger.Txn) error {
    18  		item, err := txn.Get([]byte(jwtSecret))
    19  		if err != nil {
    20  			if err == badger.ErrKeyNotFound {
    21  				return nil
    22  			}
    23  			return err
    24  		}
    25  
    26  		err = item.Value(func(val []byte) error {
    27  			secret = append([]byte{}, val...)
    28  			return nil
    29  		})
    30  		if err != nil {
    31  			return err
    32  		}
    33  		return nil
    34  	})
    35  	if err != nil {
    36  		return "", err
    37  	}
    38  
    39  	if secret == nil {
    40  		generatedJWT, err := newJWTSecret()
    41  		if err != nil {
    42  			return "", err
    43  		}
    44  		secret = []byte(generatedJWT)
    45  		err = s.main.Update(func(txn *badger.Txn) error {
    46  			return txn.SetEntry(badger.NewEntry([]byte(jwtSecret), secret))
    47  		})
    48  		if err != nil {
    49  			return "", err
    50  		}
    51  	}
    52  
    53  	return string(secret), nil
    54  }
    55  
    56  func newJWTSecret() (string, error) {
    57  	const letters = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz-"
    58  	ret := make([]byte, jwtLenght)
    59  	for i := 0; i < jwtLenght; i++ {
    60  		num, err := rand.Int(rand.Reader, big.NewInt(int64(len(letters))))
    61  		if err != nil {
    62  			return "", err
    63  		}
    64  		ret[i] = letters[num.Int64()]
    65  	}
    66  
    67  	return string(ret), nil
    68  }