From 89a78e50aa50824f0c650aaf0ab00cabfa9be20a Mon Sep 17 00:00:00 2001
From: ilija <pavlovicilija42@gmail.com>
Date: Fri, 20 Dec 2024 14:25:37 +0100
Subject: [PATCH] Improve Solana address codec modifier to handle all valid
 base58 keys

---
 pkg/solana/codec/byte_string_modifier.go      | 12 ++++++++----
 pkg/solana/codec/byte_string_modifier_test.go | 13 +++++++++++++
 pkg/solana/codec/codec_test.go                | 14 +++++++++++---
 3 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/pkg/solana/codec/byte_string_modifier.go b/pkg/solana/codec/byte_string_modifier.go
index 8a30b33d5..f33080ed4 100644
--- a/pkg/solana/codec/byte_string_modifier.go
+++ b/pkg/solana/codec/byte_string_modifier.go
@@ -22,15 +22,19 @@ func (s SolanaAddressModifier) EncodeAddress(bytes []byte) (string, error) {
 
 // DecodeAddress decodes a Base58-encoded Solana address into a 32-byte array.
 func (s SolanaAddressModifier) DecodeAddress(str string) ([]byte, error) {
-	if len(str) != 44 {
-		return nil, fmt.Errorf("%w: got length %d, expected 44 for address %s", commontypes.ErrInvalidType, len(str), str)
-	}
-
 	pubkey, err := solana.PublicKeyFromBase58(str)
 	if err != nil {
 		return nil, fmt.Errorf("%w: failed to decode Base58 address: %s", commontypes.ErrInvalidType, err)
 	}
 
+	if pubkey.IsZero() {
+		return nil, fmt.Errorf("%w: zero-value address", commontypes.ErrInvalidType)
+	}
+
+	if !pubkey.IsOnCurve() {
+		return nil, fmt.Errorf("%w: address %q with length of %d is not on the ed25519 curve", commontypes.ErrInvalidType, str, len(str))
+	}
+
 	return pubkey.Bytes(), nil
 }
 
diff --git a/pkg/solana/codec/byte_string_modifier_test.go b/pkg/solana/codec/byte_string_modifier_test.go
index 83d1bf189..e0be1ce2c 100644
--- a/pkg/solana/codec/byte_string_modifier_test.go
+++ b/pkg/solana/codec/byte_string_modifier_test.go
@@ -51,6 +51,19 @@ func TestSolanaAddressModifier(t *testing.T) {
 		assert.Contains(t, err.Error(), commontypes.ErrInvalidType.Error())
 	})
 
+	t.Run("DecodeAddress returns error for address under 32 chars", func(t *testing.T) {
+		// < than 32 chars
+		_, err := modifier.DecodeAddress("1111111111111111111111111111111")
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), commontypes.ErrInvalidType.Error())
+	})
+
+	t.Run("DecodeAddress returns error for valid length address not on the ed25519 curve", func(t *testing.T) {
+		_, err := modifier.DecodeAddress("AddressLookupTab1e11111111111111111111111111")
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), commontypes.ErrInvalidType.Error())
+	})
+
 	t.Run("Length returns 32 for Solana addresses", func(t *testing.T) {
 		assert.Equal(t, solana.PublicKeyLength, modifier.Length())
 	})
diff --git a/pkg/solana/codec/codec_test.go b/pkg/solana/codec/codec_test.go
index 8a9dedb45..689d545ce 100644
--- a/pkg/solana/codec/codec_test.go
+++ b/pkg/solana/codec/codec_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	_ "embed"
 	"slices"
+	"sync"
 	"testing"
 
 	bin "github.com/gagliardetto/binary"
@@ -52,15 +53,22 @@ func FuzzCodec(f *testing.F) {
 }
 
 type codecInterfaceTester struct {
+	accountBytesMu sync.Mutex
+	accountBytes   []byte
 	TestSelectionSupport
 }
 
 func (it *codecInterfaceTester) Setup(_ *testing.T) {}
 
 func (it *codecInterfaceTester) GetAccountBytes(_ int) []byte {
-	// TODO solana base58 string can be of variable length, this value is always 44, but it should be able to handle any length 32-44
-	pk := solana.PublicKeyFromBytes([]byte{220, 108, 195, 188, 166, 6, 163, 39, 197, 131, 44, 38, 154, 177, 232, 80, 141, 50, 7, 65, 28, 65, 182, 165, 57, 5, 176, 68, 46, 181, 58, 245})
-	return pk.Bytes()
+	it.accountBytesMu.Lock()
+	defer it.accountBytesMu.Unlock()
+	if len(it.accountBytes) != 32 {
+		pk, _ := solana.NewRandomPrivateKey()
+		it.accountBytes = pk.PublicKey().Bytes()
+	}
+
+	return it.accountBytes
 }
 
 func (it *codecInterfaceTester) GetAccountString(i int) string {