79 lines
1.7 KiB
Go
79 lines
1.7 KiB
Go
package generation
|
||
|
||
import (
|
||
cryptoRand "crypto/rand"
|
||
"encoding/binary"
|
||
"sync"
|
||
)
|
||
|
||
const (
|
||
idSize = 6 // 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
|
||
}
|