refactor(base58): replace mr-tron/base58 with in-tree long-division encoder (#425)

Author: sonicfromnewyoke

base58/base58_test.go

@@ -3,9 +3,9 @@ package base58
 import (
 	"crypto/rand"
 	"encoding/hex"
+	"fmt"
 	"testing"
 
-	mrtronbase58 "github.com/mr-tron/base58"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -98,18 +98,22 @@ func TestDecode32_Zeros(t *testing.T) {
 }
 
 func TestRoundtrip32_Random(t *testing.T) {
-	// Cross-check the specialized fixed-size path against mr-tron's
-	// well-tested general-purpose implementation.
+	// Cross-check the specialized fixed-size path against the variable-length
+	// fallback — the two share no code, so disagreement flags a bug.
 	for range 1000 {
 		var src [32]byte
 		rand.Read(src[:])
 
-		encoded := Encode32(&src)
-		assert.Equal(t, mrtronbase58.Encode(src[:]), encoded, "encode mismatch for %x", src)
+		encoded := Encode(src[:])
+		assert.Equal(t, encodeVariable(src[:]), encoded, "encode mismatch for %x", src)
 
 		var decoded [32]byte
 		require.NoError(t, Decode32(encoded, &decoded))
 		assert.Equal(t, src, decoded, "decode mismatch for %s", encoded)
+
+		generic, err := Decode(encoded)
+		require.NoError(t, err)
+		assert.Equal(t, src[:], generic, "generic decode mismatch for %s", encoded)
 	}
 }
 
@@ -118,19 +122,23 @@ func TestRoundtrip64_Random(t *testing.T) {
 		var src [64]byte
 		rand.Read(src[:])
 
-		encoded := Encode64(&src)
-		assert.Equal(t, mrtronbase58.Encode(src[:]), encoded, "encode mismatch for %x", src)
+		encoded := Encode(src[:])
+		assert.Equal(t, encodeVariable(src[:]), encoded, "encode mismatch for %x", src)
 
 		var decoded [64]byte
 		require.NoError(t, Decode64(encoded, &decoded))
 		assert.Equal(t, src, decoded, "decode mismatch for %s", encoded)
+
+		generic, err := Decode(encoded)
+		require.NoError(t, err)
+		assert.Equal(t, src[:], generic, "generic decode mismatch for %s", encoded)
 	}
 }
 
 func TestAppendEncode32_ZeroAlloc(t *testing.T) {
 	var src [32]byte
 	rand.Read(src[:])
-	expected := Encode32(&src)
+	expected := Encode(src[:])
 
 	// Pre-sized buffer: should not allocate.
 	buf := make([]byte, 0, EncodedMaxLen32)
@@ -148,7 +156,7 @@ func TestAppendEncode32_ZeroAlloc(t *testing.T) {
 func TestAppendEncode64_ZeroAlloc(t *testing.T) {
 	var src [64]byte
 	rand.Read(src[:])
-	expected := Encode64(&src)
+	expected := Encode(src[:])
 
 	buf := make([]byte, 0, EncodedMaxLen64)
 	buf = AppendEncode64(buf, &src)
@@ -162,6 +170,105 @@ func TestDecode_InvalidChars(t *testing.T) {
 	assert.Error(t, Decode32("Oinvalid", &dst)) // 'O' is not in base58
 }
 
+// Known vectors for the variable-length API. Cross-validated against
+// Bitcoin Core, bs58, and five8.
+var knownVectorsVar = []struct {
+	hex string
+	b58 string
+}{
+	{"", ""},
+	{"00", "1"},
+	{"0000", "11"},
+	{"00000000", "1111"},
+	{"61", "2g"},
+	{"626262", "a3gV"},
+	{"636363", "aPEr"},
+	{"73696d706c792061206c6f6e6720737472696e67", "2cFupjhnEsSn59qHXstmK2ffpLv2"},
+	{"00eb15231dfceb60925886b67d065299925915aeb172c06647", "1NS17iag9jJgTHD1VXjvLCEnZuQ3rJDE9L"},
+	// Solana instruction data sample from transaction_test.go.
+	{"020000003930000000000000", "3Bxs4ART6LMJ13T5"},
+}
+
+func TestEncode_KnownVectors(t *testing.T) {
+	for _, tv := range knownVectorsVar {
+		raw, err := hex.DecodeString(tv.hex)
+		require.NoError(t, err)
+		assert.Equal(t, tv.b58, Encode(raw), "hex=%s", tv.hex)
+	}
+}
+
+func TestDecode_KnownVectors(t *testing.T) {
+	for _, tv := range knownVectorsVar {
+		expected, err := hex.DecodeString(tv.hex)
+		require.NoError(t, err)
+		got, err := Decode(tv.b58)
+		require.NoError(t, err, "b58=%s", tv.b58)
+		if expected == nil {
+			expected = []byte{}
+		}
+		assert.Equal(t, expected, got, "b58=%s", tv.b58)
+	}
+}
+
+func TestEncode_Empty(t *testing.T) {
+	assert.Equal(t, "", Encode(nil))
+	assert.Equal(t, "", Encode([]byte{}))
+}
+
+func TestDecode_Empty(t *testing.T) {
+	got, err := Decode("")
+	require.NoError(t, err)
+	assert.Equal(t, []byte{}, got)
+}
+
+func TestRoundtrip_Variable_Random(t *testing.T) {
+	// Cover assorted lengths including ones the fixed-size paths can't handle.
+	for _, n := range []int{1, 5, 12, 31, 33, 63, 65, 100, 250, 1000} {
+		for range 100 {
+			src := make([]byte, n)
+			rand.Read(src)
+
+			encoded := Encode(src)
+			decoded, err := Decode(encoded)
+			require.NoError(t, err, "len=%d", n)
+			assert.Equal(t, src, decoded, "len=%d encoded=%s", n, encoded)
+		}
+	}
+}
+
+func TestRoundtrip_Variable_LeadingZeros(t *testing.T) {
+	// Encoded leading '1's must round-trip to the same number of leading zeros.
+	for zeros := 0; zeros < 10; zeros++ {
+		for tail := 0; tail < 10; tail++ {
+			src := make([]byte, zeros+tail)
+			if tail > 0 {
+				rand.Read(src[zeros:])
+				if src[zeros] == 0 {
+					src[zeros] = 1
+				}
+			}
+			encoded := Encode(src)
+			decoded, err := Decode(encoded)
+			require.NoError(t, err)
+			assert.Equal(t, src, decoded, "zeros=%d tail=%d", zeros, tail)
+		}
+	}
+}
+
+func TestDecode_InvalidChars_Variable(t *testing.T) {
+	for _, in := range []string{"0", "O", "I", "l", "abc!", "abc 123", "\x00"} {
+		_, err := Decode(in)
+		assert.Error(t, err, "expected error for %q", in)
+	}
+}
+
+func BenchmarkBase58_Decode_Variable(b *testing.B) {
+	b.SetBytes(64)
+	for b.Loop() {
+		Decode(benchStr64)
+	}
+}
+
 // Benchmarks
 var (
 	benchSrc32 [32]byte
@@ -173,8 +280,23 @@ var (
 func init() {
 	rand.Read(benchSrc32[:])
 	rand.Read(benchSrc64[:])
-	benchStr32 = Encode32(&benchSrc32)
-	benchStr64 = Encode64(&benchSrc64)
+	benchStr32 = Encode(benchSrc32[:])
+	benchStr64 = Encode(benchSrc64[:])
+}
+
+func BenchmarkBase58_EncodeVariable(b *testing.B) {
+	// Cover lengths that bypass the 32/64 fast paths and exercise the
+	// long-division encoder. Solana instruction data is typically <= 1KB.
+	for _, n := range []int{16, 100, 1000} {
+		src := make([]byte, n)
+		rand.Read(src)
+		b.Run(fmt.Sprintf("len=%d", n), func(b *testing.B) {
+			b.SetBytes(int64(n))
+			for b.Loop() {
+				Encode(src)
+			}
+		})
+	}
 }
 
 func BenchmarkBase58_Encode32(b *testing.B) {

base58/decode.go

@@ -12,6 +12,84 @@ var (
 	ErrLeadingZeros  = errors.New("base58: leading '1' count does not match leading zero bytes")
 )
 
+// Decode decodes a base58 string to bytes. Each leading '1' in encoded
+// produces a leading zero byte in the output. Empty input produces an empty
+// (non-nil) slice.
+//
+// Encoded lengths matching a 32 or 64-byte representation — the common Solana
+// sizes — are dispatched to the matrix-multiply fast paths (Decode32 /
+// Decode64), which are ~10x faster than the long-multiplication fallback. A
+// 32-byte value always encodes to 32-44 base58 chars; 64-byte to 64-88. The
+// fast paths reject inputs whose natural byte count differs (via leading-zero
+// validation), so we fall back to long multiplication on error.
+func Decode(encoded string) ([]byte, error) {
+	if len(encoded) == 0 {
+		return []byte{}, nil
+	}
+
+	encLen := len(encoded)
+	if encLen >= 32 && encLen <= EncodedMaxLen32 {
+		var dst [32]byte
+		if err := Decode32(encoded, &dst); err == nil {
+			out := make([]byte, 32)
+			copy(out, dst[:])
+			return out, nil
+		}
+	}
+	if encLen >= 64 && encLen <= EncodedMaxLen64 {
+		var dst [64]byte
+		if err := Decode64(encoded, &dst); err == nil {
+			out := make([]byte, 64)
+			copy(out, dst[:])
+			return out, nil
+		}
+	}
+
+	zeros := 0
+	for zeros < len(encoded) && encoded[zeros] == '1' {
+		zeros++
+	}
+
+	if zeros == len(encoded) {
+		return make([]byte, zeros), nil
+	}
+
+	// Upper bound on byte count of the non-leading-zero portion:
+	// ceil(n * log(58)/log(256)) ~ n * 0.7322. Use 733/1000 + 1 for safety.
+	size := ((len(encoded)-zeros)*733)/1000 + 1
+	work := make([]byte, size)
+
+	for i := zeros; i < len(encoded); i++ {
+		c := encoded[i]
+		if c < '1' || c > 'z' {
+			return nil, ErrInvalidChar
+		}
+		digit := base58Inverse[c-'1']
+		if digit == base58InvalidDigit {
+			return nil, ErrInvalidChar
+		}
+		// work = work * 58 + digit, treating work as a big-endian bigint.
+		carry := uint32(digit)
+		for j := len(work) - 1; j >= 0; j-- {
+			cur := uint32(work[j])*58 + carry
+			work[j] = byte(cur)
+			carry = cur >> 8
+		}
+		if carry != 0 {
+			return nil, ErrValueTooLarge
+		}
+	}
+
+	skip := 0
+	for skip < len(work) && work[skip] == 0 {
+		skip++
+	}
+
+	out := make([]byte, zeros+len(work)-skip)
+	copy(out[zeros:], work[skip:])
+	return out, nil
+}
+
 // Decode32 decodes a base58 string into a 32-byte array.
 func Decode32(encoded string, dst *[32]byte) error {
 	encLen := len(encoded)

base58/encode.go

@@ -5,6 +5,74 @@ import (
 	"unsafe"
 )
 
+// Encode encodes a byte slice to a base58 string. Each leading zero byte in
+// src produces a leading '1' in the output. Empty input produces an empty
+// string.
+//
+// Inputs of exactly 32 or 64 bytes — the common Solana sizes (pubkey, hash,
+// signature, private key) — are dispatched to the matrix-multiply fast paths
+// and are ~20x faster than the long-division fallback used for other lengths.
+func Encode(buf []byte) string {
+	switch len(buf) {
+	case 0:
+		return ""
+	case 32:
+		return Encode32((*[32]byte)(buf))
+	case 64:
+		return Encode64((*[64]byte)(buf))
+	default:
+		return encodeVariable(buf)
+	}
+}
+
+// encodeVariable is a long-division base58 encoder for inputs of arbitrary
+// length. Adapted from github.com/mr-tron/base58 (FastBase58Encoding); the
+// output-buffer size is corrected to zcount+size-j (upstream's
+// binsz-zcount+(size-j) panics on all-zero input and over-allocates otherwise,
+// leaving NUL bytes at the tail of the returned string).
+func encodeVariable(bin []byte) string {
+	binsz := len(bin)
+	zcount := 0
+	for zcount < binsz && bin[zcount] == 0 {
+		zcount++
+	}
+
+	// Upper bound on encoded non-zero portion: ceil(n * log(256)/log(58)) ~
+	// n * 1.366. Use 138/100 + 1 for safety.
+	size := (binsz-zcount)*138/100 + 1
+	buf := make([]byte, size)
+
+	high := size - 1
+	for i := zcount; i < binsz; i++ {
+		j := size - 1
+		for carry := uint32(bin[i]); j > high || carry != 0; j-- {
+			carry += 256 * uint32(buf[j])
+			buf[j] = byte(carry % 58)
+			carry /= 58
+			if j == 0 {
+				break
+			}
+		}
+		high = j
+	}
+
+	// Skip leading zero digits in the working buffer.
+	j := 0
+	for j < size && buf[j] == 0 {
+		j++
+	}
+
+	b58 := make([]byte, zcount+size-j)
+	for i := range zcount {
+		b58[i] = base58Chars[0]
+	}
+	for i := zcount; j < size; i++ {
+		b58[i] = base58Chars[buf[j]]
+		j++
+	}
+	return string(b58)
+}
+
 // Encode32 encodes a 32-byte array to a base58 string.
 //
 // Allocates exactly one []byte of the encoded length. For zero-allocation