diff --git a/go.mod b/go.mod index deaf8a0f..2cb3027e 100644 --- a/go.mod +++ b/go.mod @@ -93,6 +93,7 @@ require ( github.com/quic-go/qpack v0.6.0 // indirect github.com/quic-go/quic-go v0.59.1 // indirect github.com/rogpeppe/go-internal v1.14.1 // indirect + github.com/sqids/sqids-go v0.4.1 // indirect github.com/tidwall/match v1.2.0 // indirect github.com/tidwall/pretty v1.2.1 // indirect github.com/tidwall/sjson v1.2.5 // indirect diff --git a/go.sum b/go.sum index 167a2ef6..6ee6145d 100644 --- a/go.sum +++ b/go.sum @@ -177,6 +177,8 @@ github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs= github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro= github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= +github.com/sqids/sqids-go v0.4.1 h1:eQKYzmAZbLlRwHeHYPF35QhgxwZHLnlmVj9AkIj/rrw= +github.com/sqids/sqids-go v0.4.1/go.mod h1:EMwHuPQgSNFS0A49jESTfIQS+066XQTVhukrzEPScl8= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= diff --git a/pkg/idcodec/idcodec.go b/pkg/idcodec/idcodec.go new file mode 100644 index 00000000..76f94fc1 --- /dev/null +++ b/pkg/idcodec/idcodec.go @@ -0,0 +1,103 @@ +package idcodec + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/binary" + "errors" + "io" + "sync" + + "github.com/sqids/sqids-go" + "golang.org/x/crypto/hkdf" +) + +const ( + baseAlphabet = "abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ23456789" + hkdfInfo = "nezha/idcodec/alphabet/v1" + minLength = 8 + minMasterKey = 32 +) + +var ( + ErrNotInitialized = errors.New("idcodec: not initialized") + ErrInvalidCode = errors.New("idcodec: invalid id code") + ErrMasterKeyShort = errors.New("idcodec: master key too short") + + mu sync.RWMutex + encoder *sqids.Sqids +) + +func Init(masterKey []byte) error { + if len(masterKey) < minMasterKey { + return ErrMasterKeyShort + } + alphaKey := make([]byte, 32) + if _, err := io.ReadFull(hkdf.New(sha256.New, masterKey, nil, []byte(hkdfInfo)), alphaKey); err != nil { + return err + } + enc, err := sqids.New(sqids.Options{ + Alphabet: keyedShuffle(baseAlphabet, alphaKey), + MinLength: minLength, + Blocklist: []string{}, + }) + if err != nil { + return err + } + mu.Lock() + encoder = enc + mu.Unlock() + return nil +} + +func Encode(id uint64) (string, error) { + mu.RLock() + enc := encoder + mu.RUnlock() + if enc == nil { + return "", ErrNotInitialized + } + return enc.Encode([]uint64{id}) +} + +func Decode(code string) (uint64, error) { + mu.RLock() + enc := encoder + mu.RUnlock() + if enc == nil { + return 0, ErrNotInitialized + } + nums := enc.Decode(code) + if len(nums) != 1 { + return 0, ErrInvalidCode + } + if got, err := enc.Encode(nums); err != nil || got != code { + return 0, ErrInvalidCode + } + return nums[0], nil +} + +func keyedShuffle(alphabet string, key []byte) string { + runes := []rune(alphabet) + mac := hmac.New(sha256.New, key) + var counter uint64 + var pool []byte + next := func() byte { + if len(pool) == 0 { + buf := make([]byte, 8) + binary.BigEndian.PutUint64(buf, counter) + counter++ + mac.Reset() + mac.Write(buf) + pool = mac.Sum(nil) + } + b := pool[0] + pool = pool[1:] + return b + } + for i := len(runes) - 1; i > 0; i-- { + j := int(next()) % (i + 1) + runes[i], runes[j] = runes[j], runes[i] + } + return string(runes) +} diff --git a/pkg/idcodec/idcodec_test.go b/pkg/idcodec/idcodec_test.go new file mode 100644 index 00000000..0e79a991 --- /dev/null +++ b/pkg/idcodec/idcodec_test.go @@ -0,0 +1,130 @@ +package idcodec + +import ( + "strings" + "sync" + "testing" +) + +const testMasterKey = "this-is-a-32-byte-master-key-ok!" + +func resetEncoder(t *testing.T) { + t.Helper() + mu.Lock() + encoder = nil + mu.Unlock() +} + +func TestEncodeDecodeRoundTrip(t *testing.T) { + resetEncoder(t) + if err := Init([]byte(testMasterKey)); err != nil { + t.Fatalf("Init: %v", err) + } + + cases := []uint64{1, 2, 42, 1_000_000, 1<<63 - 1} + for _, id := range cases { + code, err := Encode(id) + if err != nil { + t.Fatalf("Encode(%d): %v", id, err) + } + if len(code) < minLength { + t.Fatalf("code %q shorter than min %d", code, minLength) + } + got, err := Decode(code) + if err != nil { + t.Fatalf("Decode(%q): %v", code, err) + } + if got != id { + t.Fatalf("round-trip mismatch: got %d, want %d", got, id) + } + } +} + +func TestEncodeBeforeInit(t *testing.T) { + resetEncoder(t) + if _, err := Encode(1); err != ErrNotInitialized { + t.Fatalf("Encode without Init: want ErrNotInitialized, got %v", err) + } + if _, err := Decode("abcdefgh"); err != ErrNotInitialized { + t.Fatalf("Decode without Init: want ErrNotInitialized, got %v", err) + } +} + +func TestInitRejectsShortMasterKey(t *testing.T) { + resetEncoder(t) + if err := Init([]byte("too-short")); err != ErrMasterKeyShort { + t.Fatalf("Init short master key: want ErrMasterKeyShort, got %v", err) + } +} + +func TestDecodeInvalidInputs(t *testing.T) { + resetEncoder(t) + if err := Init([]byte(testMasterKey)); err != nil { + t.Fatalf("Init: %v", err) + } + + for _, code := range []string{"", "@@@@", strings.Repeat("!", 16)} { + if _, err := Decode(code); err == nil { + t.Fatalf("Decode(%q) must fail", code) + } + } +} + +func TestAlphabetChangesWithMasterKey(t *testing.T) { + resetEncoder(t) + if err := Init([]byte(testMasterKey)); err != nil { + t.Fatalf("Init A: %v", err) + } + codeA, err := Encode(42) + if err != nil { + t.Fatalf("Encode A: %v", err) + } + + resetEncoder(t) + if err := Init([]byte(testMasterKey + "rotated-suffix-makes-key-longer!")); err != nil { + t.Fatalf("Init B: %v", err) + } + codeB, err := Encode(42) + if err != nil { + t.Fatalf("Encode B: %v", err) + } + if codeA == codeB { + t.Fatalf("rotating master key must change hashid encoding for the same id; both produced %q", codeA) + } + + if _, err := Decode(codeA); err == nil { + t.Fatalf("after rotation, old hashid %q must not decode under new key", codeA) + } +} + +func TestConcurrentEncodeDecodeIsSafe(t *testing.T) { + resetEncoder(t) + if err := Init([]byte(testMasterKey)); err != nil { + t.Fatalf("Init: %v", err) + } + var wg sync.WaitGroup + for i := 0; i < 16; i++ { + wg.Add(1) + go func(seed uint64) { + defer wg.Done() + for j := uint64(0); j < 1000; j++ { + id := seed*1000 + j + code, err := Encode(id) + if err != nil { + t.Errorf("Encode(%d): %v", id, err) + return + } + got, err := Decode(code) + if err != nil { + t.Errorf("Decode(%q): %v", code, err) + return + } + if got != id { + t.Errorf("round-trip: got %d, want %d", got, id) + return + } + } + }(uint64(i)) + } + wg.Wait() +}