Wallet v5 support for generate address

tonkeeper · May 3, 2024 · 32392d8 · 32392d8
1 parent 1dab703
commit 32392d8
Show file tree

Hide file tree

Showing 4 changed files with 78 additions and 13 deletions.
diff --git a/tonconnect/proof_test.go b/tonconnect/proof_test.go
@@ -30,6 +30,11 @@ func TestCreateSignedProof(t *testing.T) {
 			version: wallet.V4R1,
 			secret:  "another-random-secret",
 		},
+		{
+			name:    "v5r1",
+			version: wallet.V5R1,
+			secret:  "random-secret",
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -47,11 +52,11 @@ func TestCreateSignedProof(t *testing.T) {
 				t.Fatalf("SeedToPrivateKey() failed: %v", err)
 			}
 			publicKey := privateKey.Public().(ed25519.PublicKey)
-			stateInit, err := wallet.GenerateStateInit(publicKey, tt.version, 0, nil)
+			stateInit, err := wallet.GenerateStateInit(publicKey, tt.version, 0, nil, nil)
 			if err != nil {
 				t.Fatalf("GenerateStateInit() failed: %v", err)
 			}
-			accountID, err := wallet.GenerateWalletAddress(publicKey, tt.version, 0, nil)
+			accountID, err := wallet.GenerateWalletAddress(publicKey, tt.version, 0, nil, nil)
 			if err != nil {
 				t.Fatalf("GenerateWalletAddress() failed: %v", err)
 			}

diff --git a/wallet/models.go b/wallet/models.go
@@ -176,6 +176,13 @@ type DataV3 struct {
 	PublicKey   tlb.Bits256
 }
 
+type DataV4 struct {
+	Seqno       uint32
+	SubWalletId uint32
+	PublicKey   tlb.Bits256
+	PluginDict  tlb.HashmapE[tlb.Bits264, tlb.Any] // TODO: find type and check size
+}
+
 type WalletV5ID struct {
 	NetworkGlobalID uint32
 	Workchain       uint8
@@ -190,13 +197,6 @@ type DataV5 struct {
 	PluginDict tlb.HashmapE[tlb.Bits264, tlb.Uint8] // TODO: find type and check size
 }
 
-type DataV4 struct {
-	Seqno       uint32
-	SubWalletId uint32
-	PublicKey   tlb.Bits256
-	PluginDict  tlb.HashmapE[tlb.Bits264, tlb.Any] // TODO: find type and check size
-}
-
 // DataHighloadV4 represents data of a highload-wallet contract.
 type DataHighloadV4 struct {
 	SubWalletId     uint32

diff --git a/wallet/wallet.go b/wallet/wallet.go
@@ -28,7 +28,7 @@ func DefaultWalletFromSeed(seed string, blockchain blockchain) (Wallet, error) {
 // (https://github.com/toncenter/tonweb/blob/master/src/contract/wallet/WalletSources.md)
 func New(key ed25519.PrivateKey, ver Version, workchain int, subWalletId *int, blockchain blockchain) (Wallet, error) {
 	publicKey := key.Public().(ed25519.PublicKey)
-	address, err := GenerateWalletAddress(publicKey, ver, workchain, subWalletId)
+	address, err := GenerateWalletAddress(publicKey, ver, workchain, subWalletId, nil)
 	if err != nil {
 		return Wallet{}, err
 	}
@@ -56,8 +56,9 @@ func GenerateWalletAddress(
 	ver Version,
 	workchain int,
 	subWalletId *int,
+	networkGlobalID *int,
 ) (ton.AccountID, error) {
-	state, err := GenerateStateInit(key, ver, workchain, subWalletId)
+	state, err := GenerateStateInit(key, ver, workchain, subWalletId, networkGlobalID)
 	if err != nil {
 		return ton.AccountID{}, fmt.Errorf("can not generate wallet state: %v", err)
 	}
@@ -83,6 +84,7 @@ func GenerateStateInit(
 	ver Version,
 	workchain int,
 	subWalletId *int,
+	networkGlobalID *int,
 ) (tlb.StateInit, error) {
 	var (
 		err       error
@@ -112,6 +114,26 @@ func GenerateStateInit(
 			PublicKey:   publicKey,
 		}
 		err = tlb.Marshal(dataCell, data)
+	case V5R1:
+		if subWalletId == nil {
+			id := 0
+			subWalletId = &id
+		}
+		if networkGlobalID == nil {
+			id := -239 // -3 for testnet
+			networkGlobalID = &id
+		}
+		data := DataV5{
+			Seqno: 0,
+			WalletID: WalletV5ID{
+				NetworkGlobalID: uint32(*networkGlobalID),
+				Workchain:       uint8(workchain),
+				WalletVersion:   0,
+				SubWalletID:     uint32(*subWalletId),
+			},
+			PublicKey: publicKey,
+		}
+		err = tlb.Marshal(dataCell, data)
 	case HighLoadV2R2:
 		if subWalletId == nil {
 			id := DefaultSubWallet + workchain
@@ -260,7 +282,7 @@ func (w *Wallet) RawSend(
 func (w *Wallet) getInit() (tlb.StateInit, error) {
 	publicKey := w.key.Public().(ed25519.PublicKey)
 	id := int(w.subWalletId)
-	return GenerateStateInit(publicKey, w.ver, int(w.address.Workchain), &id)
+	return GenerateStateInit(publicKey, w.ver, int(w.address.Workchain), &id, nil)
 }
 
 func checkMessagesLimit(msgQty int, ver Version) error { // TODO: maybe return bool

diff --git a/wallet/wallet_test.go b/wallet/wallet_test.go
@@ -7,6 +7,7 @@ import (
 	"encoding/hex"
 	"fmt"
 	"log"
+	"reflect"
 	"testing"
 
 	"github.com/tonkeeper/tongo/boc"
@@ -55,7 +56,7 @@ func TestGenerateWalletAddress(t *testing.T) {
 	for ver, data := range testData {
 		key, _ := hex.DecodeString(data.PublicKey)
 		publicKey := ed25519.PublicKey(key)
-		address, err := GenerateWalletAddress(publicKey, ver, 0, nil)
+		address, err := GenerateWalletAddress(publicKey, ver, 0, nil, nil)
 		if err != nil {
 			t.Fatalf("address generation failed: %v", err)
 		}
@@ -169,3 +170,40 @@ func TestDeserializeMessage(t *testing.T) {
 		panic(err)
 	}
 }
+func pointer[T any](t T) *T {
+	return &t
+}
+
+func TestGenerateWalletAddress1(t *testing.T) {
+	tests := []struct {
+		name            string
+		key             ed25519.PublicKey
+		ver             Version
+		workchain       int
+		subWalletId     *int
+		networkGlobalID *int
+		want            ton.AccountID
+		wantErr         bool
+	}{
+		{
+			name:            "V5R1",
+			ver:             V5R1,
+			workchain:       0,
+			networkGlobalID: pointer(-3),
+			key:             mustPubkeyFromHex("cfa50eeb1c3293c92bd33d5aa672c1717accd8a21b96033debb6d30b5bb230df"),
+			want:            ton.MustParseAccountID("EQCsa9xhVJCw2BRL07dhxwkOoAjNHRPLN2iPggZG_ZauRYt-"),
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := GenerateWalletAddress(tt.key, tt.ver, tt.workchain, tt.subWalletId, tt.networkGlobalID)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("GenerateWalletAddress() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("GenerateWalletAddress() got = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}