Merge pull request 'Add JWT generation, creation and validation along with tests' (#1) from ProfessorFartsalot/SnowcloakUtils:main into main

Reviewed-on: Eauldane/SnowcloakUtils#1
Reviewed-by: Eauldane <elf@eauldane.com>

SnowcloakUtils.go

@@ -1,26 +0,0 @@
-package SnowcloakUtils
-
-import (
-	"crypto/rand" // Probably better than math/rand for this
-	"log"
-)
-
-func GenerateRandomString(length int) string {
-	// C# had optional parameters that allowed lowercase for chardata and gpose lobbies, Go doesn't.
-	// We can probably get away with just uppercase.
-	allowedArray := []rune("ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
-	random := make([]byte, length)
-	_, err := rand.Read(random)
-	if err != nil {
-		log.Fatalf("Couldn't generate secure slice: %s", err)
-	}
-	allowedLength := len(allowedArray)
-	i := 0
-	charArray := make([]rune, length)
-	for i < length {
-		charArray[i] = allowedArray[int(random[i])%allowedLength] // Casting random[i] due to type mismatch
-		i++
-	}
-
-	return string(charArray)
-}

files/files.go

@@ -0,0 +1,48 @@
+package files
+
+import (
+	"github.com/showwin/speedtest-go/speedtest"
+)
+
+type ConnectionStatistics struct {
+	downloadBytes uint64 // Server download speed in bytes
+	uploadBytes   uint64 // Upload speed in bytes
+	downloadSlots uint16 // How fany files can be downloaded by the server at once before queueing kicks in
+	uploadSlots   uint16 // How many uploads we can do at once
+
+	// Slots are calculated with a minimum acceptable speed of 128KiB/s per slot
+	// Using uint16 for this MIGHT get a bit munted if someone tries to use something faster than a 25gbit connection
+	// or something but I really, really doubt anyone is going to do that. Problem for later me if it turns into an issue
+}
+
+func CalculateConnectionStatistics() ConnectionStatistics {
+	var speedtestClient = speedtest.New()
+	serverList, _ := speedtestClient.FetchServers()
+	// Picks server with the lowest ping
+	targets, _ := serverList.FindServer([]int{})
+	for _, s := range targets {
+		s.PingTest(nil)
+		s.DownloadTest()
+		s.UploadTest()
+		connStats := ConnectionStatistics{
+			downloadBytes: uint64(s.DLSpeed), // Casting since speedtest-go dev implemented this as float64
+			uploadBytes:   uint64(s.ULSpeed),
+			downloadSlots: uint16(uint64(s.DLSpeed) / 131072), // 128KiB per slot
+			uploadSlots:   uint16(uint64(s.ULSpeed) / 131072), // Casting to avoid float
+		}
+
+		if connStats.downloadSlots > 1000 {
+			connStats.downloadSlots = 1000 // Max slot count should be 1000 to avoid DDoS-like behaviour
+		}
+		if connStats.uploadSlots > 1000 {
+			connStats.uploadSlots = 1000
+		}
+		s.Context.Reset()
+		return connStats
+
+	}
+	// If all else fails... fuckin' yolo it
+	connStats := ConnectionStatistics{0, 0, 1000, 1000}
+	return connStats
+
+}

files/files_test.go

@@ -0,0 +1,13 @@
+package files
+
+import (
+	"fmt"
+	"strconv"
+	"testing"
+)
+
+func TestCalculateConnectionStatistics(t *testing.T) {
+	connStats := CalculateConnectionStatistics()
+	fmt.Printf(" Download: %s bytes/s\n Upload: %s bytes/s\n Estimated Client Slots: %s\n Estimated Internal Slots: %s\n", strconv.FormatUint(connStats.downloadBytes, 10), strconv.FormatUint(connStats.uploadBytes, 10), strconv.Itoa(int(connStats.uploadSlots)), strconv.Itoa(int(connStats.downloadSlots)))
+
+}

global/global.go

@@ -0,0 +1,42 @@
+package global
+
+import (
+	crand "crypto/rand" // Probably better than math/rand for this
+	"log"
+	"math/rand"
+)
+
+type SnowcloakConfigurationBase struct {
+	Jwt string
+}
+
+func GenerateRandomString(length int) string {
+	// C# had optional parameters that allowed lowercase for chardata and gpose lobbies, Go doesn't.
+	// We can probably get away with just uppercase.
+	allowedArray := []rune("ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
+	random := make([]byte, length)
+	_, err := crand.Read(random)
+	if err != nil {
+		log.Fatalf("Couldn't generate secure slice: %s", err)
+	}
+	allowedLength := len(allowedArray)
+	i := 0
+	charArray := make([]rune, length)
+	for i < length {
+		charArray[i] = allowedArray[int(random[i])%allowedLength] // Casting random[i] due to type mismatch
+		i++
+	}
+	return string(charArray)
+}
+
+func PhasedStartupString() string {
+	// Redis is used for a lot of stuff, which means the Postgres DB gets HAMMERED on a fresh startup. Use
+	// this to generate a shuffled string that can be used to only allow e.g. UIDs starting with [x] to connect
+	// for the first however many seconds and incrementally expand access so that data doesn't need to be loaded all at
+	// once
+	chars := []rune("ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
+	rand.Shuffle(len(chars), func(i, j int) {
+		chars[i], chars[j] = chars[j], chars[i]
+	})
+	return string(chars)
+}

global/global_test.go

@@ -1,4 +1,4 @@
-package SnowcloakUtils
+package global
 
 import (
 	"fmt"
@@ -15,3 +15,7 @@ func TestGenerateRandomString(t *testing.T) {
 	fmt.Println("Test: 30 chars")
 	fmt.Println(GenerateRandomString(30))
 }
+
+func TestPhasedStartup(t *testing.T) {
+	fmt.Printf("%c", PhasedStartupString()[1])
+}

go.mod

@@ -1,3 +1,7 @@
 module SnowcloakUtils
 
 go 1.25
+
+require github.com/showwin/speedtest-go v1.7.10
+
+require github.com/golang-jwt/jwt/v5 v5.3.0 // indirect

go.sum

@@ -0,0 +1,4 @@
+github.com/golang-jwt/jwt/v5 v5.3.0 h1:pv4AsKCKKZuqlgs5sUmn4x8UlGa0kEVt/puTpKx9vvo=
+github.com/golang-jwt/jwt/v5 v5.3.0/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE=
+github.com/showwin/speedtest-go v1.7.10 h1:9o5zb7KsuzZKn+IE2//z5btLKJ870JwO6ETayUkqRFw=
+github.com/showwin/speedtest-go v1.7.10/go.mod h1:Ei7OCTmNPdWofMadzcfgq1rUO7mvJy9Jycj//G7vyfA=

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"))
+}