Add JWT generation, creation and validation along with tests

- from arieshi255/SnowcloakUtils

jwt/jwt.go

@@ -0,0 +1,114 @@
+package jwt
+
+import (
+	"encoding/base64"
+	"log"
+	"sync"
+
+	"SnowcloakUtils/global"
+
+	"github.com/golang-jwt/jwt/v5"
+)
+
+type SnowcloakClaimTypes struct {
+	Uid        string `json:"uid"`
+	Alias      string
+	CharaIdent string
+	Internal   string
+	Continent  string
+	jwt.RegisteredClaims
+}
+
+type TokenProvider struct {
+	mu     sync.Mutex
+	tokens map[string]string
+	config *global.SnowcloakConfigurationBase
+}
+
+func NewTokenProvider(cfg *global.SnowcloakConfigurationBase) *TokenProvider {
+	return &TokenProvider{
+		tokens: make(map[string]string),
+		config: cfg,
+	}
+}
+
+func (p *TokenProvider) Token() string {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	signingKey := p.config.Jwt
+
+	if token, ok := p.tokens[signingKey]; ok {
+		return token
+	}
+
+	token := p.GenerateToken("shard1", "teststringteststringteststringteststringteststringteststringteststringteststringteststringteststring") // Should be read from config file
+	p.tokens[signingKey] = token
+	return token
+}
+
+func (p *TokenProvider) GenerateToken(shard string, authSigningKey string) string {
+	token := jwt.NewWithClaims(jwt.SigningMethodHS256, SnowcloakClaimTypes{
+		Uid:      shard,
+		Internal: "true",
+	})
+
+	secret, err := base64.StdEncoding.DecodeString(authSigningKey)
+
+	if err != nil {
+		log.Fatalf("Failed to decode secret into byte array")
+	}
+
+	ss, err := token.SignedString(secret)
+
+	if err != nil {
+		log.Fatalf("Failed to sign JWT with signing key")
+	}
+
+	p.tokens[authSigningKey] = ss
+
+	// log.Printf("Generated Token: %s", ss)
+
+	return ss
+}
+
+func CreateToken(claims *SnowcloakClaimTypes, authSigningKey string) string {
+	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
+
+	secret, err := base64.StdEncoding.DecodeString(authSigningKey)
+
+	if err != nil {
+		log.Fatalf("Failed to decode secret into byte array")
+	}
+
+	ss, err := token.SignedString(secret)
+
+	if err != nil {
+		log.Fatalf("Failed to sign JWT with signing key")
+	}
+
+	return ss
+}
+
+func ValidateToken(tokenString string, authSigningKey string) jwt.MapClaims {
+	secret, err := base64.StdEncoding.DecodeString(authSigningKey)
+
+	if err != nil {
+		log.Fatalf("Failed to decode secret into byte array")
+	}
+
+	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (any, error) {
+		return secret, nil
+	}, jwt.WithValidMethods([]string{jwt.SigningMethodHS256.Alg()}))
+
+	if err != nil {
+		log.Println("Failed to parse token from string")
+		return nil
+	}
+
+	if claims, ok := token.Claims.(jwt.MapClaims); ok {
+		return claims
+	} else {
+		return nil
+	}
+}

jwt/jwt_test.go

@@ -0,0 +1,44 @@
+package jwt
+
+import (
+	"SnowcloakUtils/global"
+	"fmt"
+	"testing"
+)
+
+func TestCreateAndValidateToken(t *testing.T) {
+	fmt.Println("Test: token creation")
+	token1 := CreateToken(&SnowcloakClaimTypes{Uid: "0", CharaIdent: "test1", Alias: "myAlias1", Continent: "EU"}, "teststringteststringteststringteststringteststringteststringteststringteststringteststringteststring")
+	fmt.Println(token1)
+	fmt.Println("Test: token validation")
+	fmt.Println(ValidateToken(token1, "teststringteststringteststringteststringteststringteststringteststringteststringteststringteststring"))
+	fmt.Println("Test: bad token creation")
+	token2 := CreateToken(&SnowcloakClaimTypes{Uid: "423235", CharaIdent: "test1", Alias: "myAlias", Continent: "EU"}, "teststringteststringteststringteststringteststringteststringteststringteststringteststringteststring")
+	fmt.Println(token2)
+	fmt.Println("Test: bad token validation")
+	badValidate := ValidateToken(token2, "dGVzdHRlc3Rpbmd0dGVzdGluZw==")
+	if badValidate != nil {
+		t.Errorf("Token wasn't invalid")
+	} else {
+		fmt.Println("Didn't validate invalid token")
+	}
+}
+
+func TestGenerateToken(t *testing.T) {
+	tokenProvider := NewTokenProvider(&global.SnowcloakConfigurationBase{Jwt: "teststringteststringteststringteststringteststringteststringteststringteststringteststringteststring"})
+
+	fmt.Println("Test: token generation")
+	fmt.Println(tokenProvider.GenerateToken("shard2", "teststringteststringteststringteststringteststringteststringteststringteststringteststringteststring"))
+	fmt.Println("Test: get cached token")
+	fmt.Println(tokenProvider.Token())
+	fmt.Println("Test: update token")
+	fmt.Println(tokenProvider.GenerateToken("shard3", "teststringteststringteststringteststringteststringteststringteststringteststringteststringteststring"))
+	fmt.Println("Test: get updated token")
+	fmt.Println(tokenProvider.Token())
+	fmt.Println("Test: create bad token")
+	fmt.Println(tokenProvider.GenerateToken("shard4", "dGVzdHRlc3Rpbmd0dGVzdGluZw=="))
+	fmt.Println("Test: get bad token")
+	fmt.Println(tokenProvider.Token())
+	fmt.Println("Test: validate cached token")
+	fmt.Println(ValidateToken(tokenProvider.Token(), "teststringteststringteststringteststringteststringteststringteststringteststringteststringteststring"))
+}