Secure Token Generator — Go source

// Package tokengenerator is the Go twin of CosmoDev's src/lib/token-generator.ts
// (dual source: the web lib is TypeScript, the CLI lib is Go — kept in
// lock-step). It provides secure, unbiased token generation with an injectable
// RNG so generation is unit-testable with a seeded PRNG for exact, reproducible
// outputs. Pure + deterministic, never panics — invalid input returns "" or 0,
// mirroring the TS lib. The table-driven tests in token-generator_test.go share
// vectors with src/lib/token-generator.test.ts so the two implementations are
// held to the same contract.
//
// Algorithm mirrors the TS lib exactly: resolve the effective alphabet (explicit
// alphabet → encoding → hex), compute the output length that carries the
// requested entropy, then sample each symbol without modulo bias via rejection
// sampling so non-power-of-two alphabets (base62, alphanumeric) stay uniform.
package tokengenerator

import (
	cryptorand "crypto/rand"
	"encoding/binary"
	"math"
)

// RNG is a uniform float generator in [0, 1), the same shape as the TS lib's
// `rng?: () => number`. The default (DefaultRNG) is a CSPRNG; tests inject a
// seeded generator (mulberry32) for reproducible output.
type RNG func() float64

// Alphabet is the output character set. The zero value ("") means "unset", which
// makes ResolveAlphabet fall through to the encoding, then to hex — matching the
// TS lib's optional `alphabet` field. AlphabetCustom reads its symbols from
// Options.CustomAlphabet.
type Alphabet string

const (
	AlphabetHex             Alphabet = "hex"
	AlphabetHexUpper        Alphabet = "hex-upper"
	AlphabetBase32          Alphabet = "base32"
	AlphabetBase32Crockford Alphabet = "base32-crockford"
	AlphabetBase64          Alphabet = "base64"
	AlphabetBase64URL       Alphabet = "base64url"
	AlphabetBase62          Alphabet = "base62"
	AlphabetAlphanumeric    Alphabet = "alphanumeric"
	AlphabetCustom          Alphabet = "custom"
)

// Encoding is a convenience alias that maps 1:1 onto an Alphabet (the common
// encodings). The zero value ("") means "unset".
type Encoding string

const (
	EncodingHex          Encoding = "hex"
	EncodingBase32       Encoding = "base32"
	EncodingBase64       Encoding = "base64"
	EncodingBase64URL    Encoding = "base64url"
	EncodingBase62       Encoding = "base62"
	EncodingAlphanumeric Encoding = "alphanumeric"
)

// Alphabets holds the fixed alphabet strings. It is the Go twin of the exported
// ALPHABETS record in src/lib/token-generator.ts. Custom is intentionally absent
// — it is caller-supplied via Options.CustomAlphabet.
var Alphabets = map[Alphabet]string{
	AlphabetHex:             "0123456789abcdef",
	AlphabetHexUpper:        "0123456789ABCDEF",
	AlphabetBase32:          "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567",                // RFC 4648
	AlphabetBase32Crockford: "0123456789ABCDEFGHJKMNPQRSTVWXYZ",               // Crockford (no I/L/O/U)
	AlphabetBase64:          "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/",
	AlphabetBase64URL:       "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_",
	AlphabetBase62:          "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz",
	AlphabetAlphanumeric:    "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ",
}

// encodingToAlphabet is the unexported mirror of ENCODING_TO_ALPHABET (not
// exported in the TS lib either).
var encodingToAlphabet = map[Encoding]Alphabet{
	EncodingHex:          AlphabetHex,
	EncodingBase32:       AlphabetBase32,
	EncodingBase64:       AlphabetBase64,
	EncodingBase64URL:    AlphabetBase64URL,
	EncodingBase62:       AlphabetBase62,
	EncodingAlphanumeric: AlphabetAlphanumeric,
}

// Options configures GenerateToken. It is the Go twin of GenerateOptions. The
// zero value (Options{}) — or a nil *Options — resolves to the hex alphabet with
// the default CSPRNG, matching `generateToken(bytes)` with no options in TS.
//
// Alphabet wins over Encoding when both are set (Alphabet != ""). AlphabetCustom
// makes the generator read symbols from CustomAlphabet. RNG, when nil, uses
// DefaultRNG (the CSPRNG).
type Options struct {
	Alphabet       Alphabet // "" (zero) → fall through to Encoding → hex
	CustomAlphabet string   // used only when Alphabet == AlphabetCustom
	RNG            RNG      // nil → DefaultRNG (CSPRNG)
}

// DefaultRNG is the CSPRNG-backed [0, 1) float used when no RNG is injected. It
// is the Go twin of the unexported defaultRng in src/lib/token-generator.ts: a
// 32-bit crypto draw divided by 2^32.
func DefaultRNG() float64 {
	var buf [4]byte
	if _, err := cryptorand.Read(buf[:]); err != nil {
		return 0 // crypto/rand essentially never fails; on failure, fall to 0.
	}
	v := binary.BigEndian.Uint32(buf[:])
	return float64(v) / 4294967296.0 // / 0x100000000
}

// ResolveAlphabet resolves the effective alphabet string from options and an
// optional encoding. It is the Go twin of resolveAlphabet.
//
// Precedence: explicit Alphabet → Encoding → hex. AlphabetCustom with no/empty
// CustomAlphabet resolves to "" (an invalid, empty set).
func ResolveAlphabet(opts *Options, encoding Encoding) string {
	if opts != nil && opts.Alphabet == AlphabetCustom {
		return opts.CustomAlphabet
	}
	if opts != nil && opts.Alphabet != "" {
		return Alphabets[opts.Alphabet]
	}
	if encoding != "" {
		return Alphabets[encodingToAlphabet[encoding]]
	}
	return Alphabets[AlphabetHex]
}

// ConstantTimeSelect builds an n-character string from alphabet, selecting each
// symbol without modulo bias via rejection sampling: a 32-bit draw is rejected
// if it falls in the uneven remainder, so every symbol stays equally likely —
// important for non-power-of-two alphabets like base62. It is the Go twin of
// constantTimeSelect. A nil rng uses DefaultRNG. It returns "" for an empty
// alphabet or non-positive n.
//
// The rejection cap (guard) keeps a pathological rng from looping forever.
func ConstantTimeSelect(alphabet string, n int, rng RNG) string {
	size := len(alphabet)
	if size < 1 || n < 1 {
		return ""
	}
	if rng == nil {
		rng = DefaultRNG
	}
	// Largest multiple of size ≤ 2^32-1; draws at/above it are re-rolled.
	limit := math.Floor(4294967295.0/float64(size)) * float64(size)
	out := make([]byte, 0, n)
	for i := 0; i < n; i++ {
		x := rng() * 4294967296.0 // [0, 2^32)
		guard := 0
		for x >= limit && guard < 64 {
			x = rng() * 4294967296.0
			guard++
		}
		idx := int(int64(math.Floor(x)) % int64(size))
		out = append(out, alphabet[idx])
	}
	return string(out)
}

// OutputLength returns the number of output characters needed to carry `bytes`
// bytes of entropy in an alphabet of `alphabetSize` symbols. It is the Go twin
// of outputLength. Returns 0 for non-positive bytes or an alphabet smaller than
// 2 symbols.
func OutputLength(bytes, alphabetSize int) int {
	if bytes < 1 || alphabetSize < 2 {
		return 0
	}
	return int(math.Ceil((float64(bytes) * 8) / math.Log2(float64(alphabetSize))))
}

// GenerateToken generates a token with `bytes` bytes of underlying entropy,
// rendered through opts.Alphabet (or encoding). Each character is sampled
// uniformly without modulo bias, so the output is unbiased even for
// base62/alphanumeric. It is the Go twin of generateToken. Returns "" for
// invalid input (non-positive bytes, an alphabet with fewer than 2 symbols).
//
// Example: GenerateToken(16, &Options{Alphabet: AlphabetHex}, "") → 32 hex
// chars (128 bits).
func GenerateToken(bytes int, opts *Options, encoding Encoding) string {
	if bytes < 1 {
		return ""
	}
	alphabet := ResolveAlphabet(opts, encoding)
	if len(alphabet) < 2 {
		return ""
	}
	n := OutputLength(bytes, len(alphabet))
	rng := DefaultRNG
	if opts != nil && opts.RNG != nil {
		rng = opts.RNG
	}
	return ConstantTimeSelect(alphabet, n, rng)
}

// EstimateEntropy returns the entropy (in bits) of a token of `bytes` entropy in
// an alphabet of `alphabetSize` symbols. It equals OutputLength * log2(size),
// which is ≥ bytes*8 because the char count is rounded up. It is the Go twin of
// estimateEntropy. Returns 0 for invalid input.
func EstimateEntropy(bytes, alphabetSize int) float64 {
	if bytes < 1 || alphabetSize < 2 {
		return 0
	}
	return float64(OutputLength(bytes, alphabetSize)) * math.Log2(float64(alphabetSize))
}

// StrengthLabel buckets an entropy estimate (bits) into a human strength label.
// It is the Go twin of strengthLabel. Tiers: weak <64 · fair 64–127 · strong
// 128–255 · very strong ≥256. Non-finite input (NaN/Inf) is treated as weak, the
// conservative default.
func StrengthLabel(entropyBits float64) string {
	if math.IsNaN(entropyBits) || math.IsInf(entropyBits, 0) || entropyBits < 64 {
		return "weak"
	}
	if entropyBits < 128 {
		return "fair"
	}
	if entropyBits < 256 {
		return "strong"
	}
	return "very strong"
}