diff options
| author | Ethan Werner <mail@ethanwerner.com> | 2026-08-07 17:17:32 +0400 |
|---|---|---|
| committer | Ethan Werner <mail@ethanwerner.com> | 2026-08-07 17:17:32 +0400 |
| commit | e4e84c7aad9fb887150ba1fcfaf64463765419c9 (patch) | |
| tree | cdafa934aa61d8b646435d2a7b4ed0d7ac5c9f55 /go | |
| parent | e0f9cc69f9f7f524610d65b288eecdbf2ad958d6 (diff) | |
Diffstat (limited to 'go')
| -rw-r--r-- | go/snowflake.go | 196 | ||||
| -rw-r--r-- | go/snowflake_bench_test.go | 65 | ||||
| -rw-r--r-- | go/snowflake_fuzz_test.go | 102 | ||||
| -rw-r--r-- | go/snowflake_test.go | 340 |
4 files changed, 703 insertions, 0 deletions
diff --git a/go/snowflake.go b/go/snowflake.go new file mode 100644 index 0000000..92a6124 --- /dev/null +++ b/go/snowflake.go @@ -0,0 +1,196 @@ +// Package snowflake generates random, URL-safe identifiers. +// +// Each ID is 16 bytes of cryptographically secure randomness, encoded as a +// base62 string. IDs implement encoding.TextMarshaler, json.Marshaler, +// and database/sql.Scanner/Valuer for convenient serialization. +package snowflake + +import ( + "crypto/rand" + "database/sql/driver" + "encoding/hex" + "errors" + "fmt" +) + +const ( + // IDByteSize is the size of a Snowflake ID in bytes. + IDByteSize = 16 + // Alphabet is the base62 character set used for encoding. + Alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" +) + +var ( + ErrInvalidIDLength = errors.New("snowflake: invalid ID byte length") + ErrInvalidEncoding = errors.New("snowflake: invalid encoding string") + decodeLUT [256]int8 +) + +func init() { + for i := range decodeLUT { + decodeLUT[i] = -1 + } + + for i := 0; i < len(Alphabet); i++ { + decodeLUT[Alphabet[i]] = int8(i) + } +} + +// Snowflake is a 16-byte random identifier. +type Snowflake struct { + bytes [IDByteSize]byte +} + +// New returns a new random Snowflake ID. +func New() Snowflake { + var s Snowflake + if _, err := rand.Read(s.bytes[:]); err != nil { + panic(err) + } + return s +} + +// Zero returns the zero value of Snowflake. +func Zero() Snowflake { + return Snowflake{} +} + +// IsZero reports whether s is the zero value. +func (s Snowflake) IsZero() bool { + return s == Snowflake{} +} + +// String returns the base62 encoding of the ID. +func (s Snowflake) String() string { + base := uint64(len(Alphabet)) + var res [64]byte + i := len(res) + + tmp := s.bytes + for { + var remainder uint64 + allZero := true + for j := 0; j < IDByteSize; j++ { + accumulator := uint64(tmp[j]) + (remainder << 8) + tmp[j] = byte(accumulator / base) + remainder = accumulator % base + if tmp[j] > 0 { + allZero = false + } + } + i-- + res[i] = Alphabet[remainder] + if allZero { + break + } + } + return string(res[i:]) +} + +// FromString parses a base62-encoded Snowflake ID. +func FromString(s string) (Snowflake, error) { + var res Snowflake + base := uint64(len(Alphabet)) + + for i := 0; i < len(s); i++ { + charVal := decodeLUT[s[i]] + if charVal < 0 { + return res, ErrInvalidEncoding + } + + carry := uint64(charVal) + for j := IDByteSize - 1; j >= 0; j-- { + accumulator := uint64(res.bytes[j])*base + carry + res.bytes[j] = byte(accumulator) + carry = accumulator >> 8 + } + + if carry > 0 { + return Snowflake{}, ErrInvalidEncoding + } + } + return res, nil +} + +// Hex returns the hex-encoded ID. +func (s Snowflake) Hex() string { + return hex.EncodeToString(s.bytes[:]) +} + +// FromHex parses a hex-encoded Snowflake ID. +func FromHex(h string) (Snowflake, error) { + decoded, err := hex.DecodeString(h) + if err != nil { + return Snowflake{}, err + } + if len(decoded) != IDByteSize { + return Snowflake{}, ErrInvalidIDLength + } + var s Snowflake + copy(s.bytes[:], decoded) + return s, nil +} + +// MarshalText implements encoding.TextMarshaler. +func (s Snowflake) MarshalText() ([]byte, error) { + return []byte(s.String()), nil +} + +// UnmarshalText implements encoding.TextUnmarshaler. +func (s *Snowflake) UnmarshalText(text []byte) error { + id, err := FromString(string(text)) + if err != nil { + return err + } + *s = id + return nil +} + +// MarshalJSON implements json.Marshaler. +func (s Snowflake) MarshalJSON() ([]byte, error) { + str := s.String() + res := make([]byte, 0, len(str)+2) + res = append(res, '"') + res = append(res, str...) + res = append(res, '"') + return res, nil +} + +// UnmarshalJSON implements json.Unmarshaler. +func (s *Snowflake) UnmarshalJSON(data []byte) error { + if len(data) < 2 || data[0] != '"' || data[len(data)-1] != '"' { + return ErrInvalidEncoding + } + return s.UnmarshalText(data[1 : len(data)-1]) +} + +// Value implements driver.Valuer for database/sql support. +func (s Snowflake) Value() (driver.Value, error) { + b := make([]byte, IDByteSize) + copy(b, s.bytes[:]) + return b, nil +} + +// Scan implements sql.Scanner for database/sql support. +func (s *Snowflake) Scan(value interface{}) error { + if value == nil { + s.bytes = [IDByteSize]byte{} + return nil + } + + var source []byte + switch t := value.(type) { + case []byte: + source = t + case string: + source = []byte(t) + default: + return fmt.Errorf("snowflake: cannot scan type %T", value) + } + + if len(source) != IDByteSize { + return ErrInvalidIDLength + } + copy(s.bytes[:], source) + return nil +} diff --git a/go/snowflake_bench_test.go b/go/snowflake_bench_test.go new file mode 100644 index 0000000..ffad399 --- /dev/null +++ b/go/snowflake_bench_test.go @@ -0,0 +1,65 @@ +package snowflake + +import "testing" + +var benchmarkSnowflake = New() + +func BenchmarkNew(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + New() + } +} + +func BenchmarkSnowflake_String(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + benchmarkSnowflake.String() + } +} + +func BenchmarkFromString(b *testing.B) { + s := benchmarkSnowflake.String() + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + FromString(s) + } +} + +func BenchmarkSnowflake_Hex(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + benchmarkSnowflake.Hex() + } +} + +func BenchmarkFromHex(b *testing.B) { + h := benchmarkSnowflake.Hex() + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + FromHex(h) + } +} + +func BenchmarkMarshalJSON(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + benchmarkSnowflake.MarshalJSON() + } +} + +func BenchmarkMarshalText(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + benchmarkSnowflake.MarshalText() + } +} + +func BenchmarkValue(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + benchmarkSnowflake.Value() + } +} diff --git a/go/snowflake_fuzz_test.go b/go/snowflake_fuzz_test.go new file mode 100644 index 0000000..a2ba624 --- /dev/null +++ b/go/snowflake_fuzz_test.go @@ -0,0 +1,102 @@ +package snowflake + +import "testing" + +func FuzzFromString(f *testing.F) { + s := New() + f.Add(s.String()) + f.Add("") + f.Add("0") + f.Add("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789") + + f.Fuzz(func(t *testing.T, data string) { + result, err := FromString(data) + if err != nil { + return + } + if !result.IsZero() { + roundTrip := result.String() + parsed, err := FromString(roundTrip) + if err != nil { + t.Fatalf("round-trip failed: %v", err) + } + if parsed != result { + t.Fatalf("round-trip mismatch: %v != %v", result, parsed) + } + } + }) +} + +func FuzzFromHex(f *testing.F) { + s := New() + f.Add(s.Hex()) + f.Add("") + f.Add("00") + + f.Fuzz(func(t *testing.T, data string) { + result, err := FromHex(data) + if err != nil { + return + } + roundTrip := result.Hex() + parsed, err := FromHex(roundTrip) + if err != nil { + t.Fatalf("round-trip failed: %v", err) + } + if parsed != result { + t.Fatalf("round-trip mismatch: %v != %v", result, parsed) + } + }) +} + +func FuzzUnmarshalJSON(f *testing.F) { + s := New() + serialized, _ := s.MarshalJSON() + f.Add(string(serialized)) + f.Add(`""`) + f.Add(`"0"`) + + f.Fuzz(func(t *testing.T, data string) { + var result Snowflake + err := result.UnmarshalJSON([]byte(data)) + if err != nil { + return + } + if !result.IsZero() { + roundTrip, _ := result.MarshalJSON() + var parsed Snowflake + if err := parsed.UnmarshalJSON(roundTrip); err != nil { + t.Fatalf("round-trip failed: %v", err) + } + if parsed != result { + t.Fatalf("round-trip mismatch: %v != %v", result, parsed) + } + } + }) +} + +func FuzzUnmarshalText(f *testing.F) { + s := New() + serialized, _ := s.MarshalText() + f.Add(string(serialized)) + f.Add("") + f.Add("0") + + f.Fuzz(func(t *testing.T, data string) { + var result Snowflake + err := result.UnmarshalText([]byte(data)) + if err != nil { + return + } + if !result.IsZero() { + roundTrip, _ := result.MarshalText() + var parsed Snowflake + if err := parsed.UnmarshalText(roundTrip); err != nil { + t.Fatalf("round-trip failed: %v", err) + } + if parsed != result { + t.Fatalf("round-trip mismatch: %v != %v", result, parsed) + } + } + }) +} diff --git a/go/snowflake_test.go b/go/snowflake_test.go new file mode 100644 index 0000000..e69644c --- /dev/null +++ b/go/snowflake_test.go @@ -0,0 +1,340 @@ +package snowflake + +import ( + "encoding/json" + "testing" +) + +func TestNew(t *testing.T) { + s := New() + if s.IsZero() { + t.Error("New() returned zero Snowflake") + } +} + +func TestNewUnique(t *testing.T) { + seen := make(map[string]bool, 1000) + + for i := 0; i < 1000; i++ { + s := New() + + str := s.String() + if seen[str] { + t.Fatalf("duplicate ID generated: %s", str) + } + seen[str] = true + } +} + +func TestIsZero(t *testing.T) { + zero := Zero() + if !zero.IsZero() { + t.Error("Zero Snowflake should be zero") + } + + nonZero := New() + if nonZero.IsZero() { + t.Error("New Snowflake should not be zero") + } +} + +func TestStringRoundTrip(t *testing.T) { + original := New() + str := original.String() + parsed, err := FromString(str) + if err != nil { + t.Fatalf("FromString failed: %v", err) + } + + if parsed != original { + t.Errorf("round-trip mismatch: original=%+v, parsed=%+v", original, parsed) + } +} + +func TestStringRoundTripMany(t *testing.T) { + for i := 0; i < 100; i++ { + original := New() + + str := original.String() + parsed, err := FromString(str) + if err != nil { + t.Fatalf("iteration %d: FromString failed: %v", i, err) + } + if parsed != original { + t.Fatalf("iteration %d: round-trip mismatch", i) + } + } +} + +func TestFromStringInvalidChars(t *testing.T) { + _, err := FromString("!!invalid!!") + if err != ErrInvalidEncoding { + t.Errorf("expected ErrInvalidEncoding, got: %v", err) + } +} + +func TestFromStringEmpty(t *testing.T) { + s, err := FromString("") + if err != nil { + t.Errorf("unexpected error for empty string: %v", err) + } + if !s.IsZero() { + t.Error("empty string should produce zero Snowflake") + } +} + +func TestFromStringTooLong(t *testing.T) { + _, err := FromString("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789") + if err != ErrInvalidEncoding { + t.Errorf("expected ErrInvalidEncoding for overflow, got: %v", err) + } +} + +func TestHexString(t *testing.T) { + s := New() + hex := s.Hex() + + if len(hex) != IDByteSize*2 { + t.Errorf("expected hex length %d, got %d", IDByteSize*2, len(hex)) + } + + for _, c := range hex { + if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f')) { + t.Errorf("invalid hex character: %c", c) + } + } +} + +func TestZeroString(t *testing.T) { + zero := Zero() + str := zero.String() + if str != "0" { + t.Errorf("Zero Snowflake should stringify to %q, got: %q", "0", str) + } + + parsed, err := FromString(str) + if err != nil { + t.Fatalf("FromString(%q) failed: %v", str, err) + } + if !parsed.IsZero() { + t.Error("FromString(Zero().String()) should round-trip to zero") + } +} + +func TestHexRoundTrip(t *testing.T) { + original := New() + hex := original.Hex() + parsed, err := FromHex(hex) + if err != nil { + t.Fatalf("FromHex failed: %v", err) + } + if parsed != original { + t.Error("hex round-trip mismatch") + } +} + +func TestFromHexInvalid(t *testing.T) { + _, err := FromHex("xyz") + if err == nil { + t.Error("expected error for invalid hex") + } +} + +func TestFromHexWrongLength(t *testing.T) { + _, err := FromHex("01") + if err != ErrInvalidIDLength { + t.Errorf("expected ErrInvalidIDLength, got: %v", err) + } +} + +func TestMarshalText(t *testing.T) { + s := New() + text, err := s.MarshalText() + if err != nil { + t.Fatalf("MarshalText failed: %v", err) + } + + var decoded Snowflake + err = decoded.UnmarshalText(text) + if err != nil { + t.Fatalf("UnmarshalText failed: %v", err) + } + + if decoded != s { + t.Error("text marshal/unmarshal mismatch") + } +} + +func TestUnmarshalTextInvalid(t *testing.T) { + var s Snowflake + err := s.UnmarshalText([]byte("!!!")) + if err == nil { + t.Error("expected error for invalid text") + } +} + +func TestMarshalJSON(t *testing.T) { + s := New() + data, err := s.MarshalJSON() + if err != nil { + t.Fatalf("MarshalJSON failed: %v", err) + } + + var decoded Snowflake + err = json.Unmarshal(data, &decoded) + if err != nil { + t.Fatalf("json.Unmarshal failed: %v", err) + } + + if decoded != s { + t.Error("JSON marshal/unmarshal mismatch") + } +} + +func TestUnmarshalJSONMissingQuotes(t *testing.T) { + var s Snowflake + err := s.UnmarshalJSON([]byte(`invalid`)) + if err != ErrInvalidEncoding { + t.Errorf("expected ErrInvalidEncoding, got: %v", err) + } +} + +func TestUnmarshalJSONEmpty(t *testing.T) { + var s Snowflake + err := s.UnmarshalJSON([]byte(`""`)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !s.IsZero() { + t.Error("empty JSON string should produce zero Snowflake") + } +} + +func TestUnmarshalJSONInvalidInner(t *testing.T) { + var s Snowflake + err := s.UnmarshalJSON([]byte(`"!!!"`)) + if err != ErrInvalidEncoding { + t.Errorf("expected ErrInvalidEncoding for invalid inner content, got: %v", err) + } +} + +func TestDriverValue(t *testing.T) { + s := New() + val, err := s.Value() + if err != nil { + t.Fatalf("Value failed: %v", err) + } + + bytes, ok := val.([]byte) + if !ok { + t.Fatal("Value should return []byte") + } + + if len(bytes) != IDByteSize { + t.Errorf("expected byte length %d, got %d", IDByteSize, len(bytes)) + } +} + +func TestDriverScan(t *testing.T) { + original := New() + bytes := make([]byte, IDByteSize) + copy(bytes, original.bytes[:]) + + var s Snowflake + if err := s.Scan(bytes); err != nil { + t.Fatalf("Scan failed: %v", err) + } + + if s != original { + t.Error("Scan produced mismatched value") + } +} + +func TestDriverScanString(t *testing.T) { + original := New() + bytes := make([]byte, IDByteSize) + copy(bytes, original.bytes[:]) + + var s Snowflake + if err := s.Scan(string(bytes)); err != nil { + t.Fatalf("Scan(string) failed: %v", err) + } + + if s != original { + t.Error("Scan(string) produced mismatched value") + } +} + +func TestDriverScanNil(t *testing.T) { + var s Snowflake + err := s.Scan(nil) + if err != nil { + t.Fatalf("Scan(nil) failed: %v", err) + } + + if !s.IsZero() { + t.Error("Scan(nil) should produce zero Snowflake") + } +} + +func TestDriverScanWrongLength(t *testing.T) { + var s Snowflake + err := s.Scan([]byte{1, 2, 3}) + if err != ErrInvalidIDLength { + t.Errorf("expected ErrInvalidIDLength, got: %v", err) + } +} + +func TestDriverScanZeroBytes(t *testing.T) { + var s Snowflake + zeroBytes := make([]byte, IDByteSize) + err := s.Scan(zeroBytes) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !s.IsZero() { + t.Error("Scan with all-zero bytes should produce zero Snowflake") + } +} + +func TestDriverScanUnsupportedType(t *testing.T) { + var s Snowflake + err := s.Scan(42) + if err == nil { + t.Error("expected error for unsupported type") + } +} + +func TestJSONEmbedding(t *testing.T) { + type Wrapper struct { + Name string `json:"name"` + ID Snowflake `json:"id"` + Ref *Snowflake `json:"ref,omitempty"` + } + + s := New() + w := Wrapper{ + Name: "test", + ID: s, + Ref: &s, + } + + data, err := json.Marshal(&w) + if err != nil { + t.Fatalf("marshal failed: %v", err) + } + + var decoded Wrapper + err = json.Unmarshal(data, &decoded) + if err != nil { + t.Fatalf("unmarshal failed: %v", err) + } + + if decoded.ID != s { + t.Error("embedded ID mismatch") + } + + if decoded.Ref == nil || *decoded.Ref != s { + t.Error("embedded ref mismatch") + } +} |