up

pkg/generation/base62.go

@@ -0,0 +1,62 @@
+package generation
+
+import (
+	"errors"
+)
+
+const alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+
+var (
+	base = len(alphabet)
+	rev  [256]int
+)
+
+func init() {
+	for i := range rev {
+		rev[i] = -1
+	}
+	for i, c := range alphabet {
+		rev[c] = i
+	}
+}
+
+// EncodeBase62 encodes a byte slice into a base62 string
+func EncodeBase62(b []byte) string {
+	var n uint64
+	for i := 0; i < len(b); i++ {
+		n = n<<8 + uint64(b[i])
+	}
+
+	if n == 0 {
+		return string(alphabet[0])
+	}
+
+	var s []byte
+	for n > 0 {
+		r := n % uint64(base)
+		s = append([]byte{alphabet[r]}, s...)
+		n = n / uint64(base)
+	}
+	return string(s)
+}
+
+// DecodeBase62 decodes a base62 string into a byte slice
+func DecodeBase62(s string) ([]byte, error) {
+	var n uint64
+	for i := 0; i < len(s); i++ {
+		c := s[i]
+		val := rev[c]
+		if val == -1 {
+			return nil, errors.New("invalid base62 character")
+		}
+		n = n*uint64(base) + uint64(val)
+	}
+
+	// Recover byte slice from integer
+	b := make([]byte, rawIDLength)
+	for i := rawIDLength - 1; i >= 0; i-- {
+		b[i] = byte(n & 0xff)
+		n >>= 8
+	}
+	return b, nil
+}

pkg/generation/generation.go

@@ -0,0 +1,78 @@
+package generation
+
+import (
+	cryptoRand "crypto/rand"
+	"encoding/binary"
+	"sync"
+)
+
+const (
+	idSize      = 4                   // 4 bytes for random part
+	prefixSize  = 2                   // 2 bytes for client prefix
+	rawIDLength = prefixSize + idSize // total 6 bytes
+	base62Len   = 8                   // 6 bytes encoded in base62 ~ 8 chars
+	poolSize    = 10000
+)
+
+type Generator struct {
+	prefix  [2]byte
+	idPool  chan string
+	mu      sync.Mutex
+	started bool
+}
+
+// NewGenerator initializes the generator with a 2-byte prefix
+func NewGenerator(prefix uint16) *Generator {
+	var b [2]byte
+	binary.BigEndian.PutUint16(b[:], prefix)
+
+	g := &Generator{
+		prefix: b,
+		idPool: make(chan string, poolSize),
+	}
+	go g.fillPool()
+	return g
+}
+
+// fillPool pre-generates short IDs into the pool
+func (g *Generator) fillPool() {
+	for {
+		for i := 0; i < poolSize/10; i++ {
+			id := g.generateRawID()
+			g.idPool <- EncodeBase62(id)
+		}
+	}
+}
+
+// generateRawID creates 6 bytes: 2-byte prefix + 4-byte random
+func (g *Generator) generateRawID() []byte {
+	random := make([]byte, idSize)
+	_, err := cryptoRand.Read(random)
+	if err != nil {
+		panic("failed to read random bytes: " + err.Error())
+	}
+
+	raw := make([]byte, rawIDLength)
+	copy(raw[0:2], g.prefix[:])
+	copy(raw[2:], random)
+	return raw
+}
+
+// NextID returns the next available short ID from the pool
+func (g *Generator) NextID() string {
+	return <-g.idPool
+}
+
+// DecodeID (optional) – for analytics/debugging
+func DecodeID(shortID string) (prefix uint16, randomPart []byte, err error) {
+	raw, err := DecodeBase62(shortID)
+	if err != nil {
+		return 0, nil, err
+	}
+	if len(raw) != rawIDLength {
+		return 0, nil, err
+	}
+	prefix = binary.BigEndian.Uint16(raw[0:2])
+	randomPart = raw[2:]
+	return
+}