Support versioned transactions

Author: Jeff Yanta

pkg/jupiter/client.go

@@ -51,6 +51,10 @@ func (q *Quote) GetEstimatedSwapAmount() uint64 {
 	return q.estimatedSwapAmount
 }
 
+func (q *Quote) String() string {
+	return q.jsonString
+}
+
 // GetQuote gets an optimal route for performing a swap
 func (c *Client) GetQuote(
 	ctx context.Context,
@@ -117,6 +121,7 @@ type SwapInstructions struct {
 	SetupInstructions         []solana.Instruction
 	SwapInstruction           solana.Instruction
 	CleanupInstruction        *solana.Instruction
+	AddressLookupTables       []string
 }
 
 // GetSwapInstructions gets the instructions to construct a transaction to sign
@@ -130,10 +135,6 @@ func (c *Client) GetSwapInstructions(
 	tracer := metrics.TraceMethodCall(ctx, metricsStructName, "GetSwapInstructions")
 	defer tracer.End()
 
-	if !quote.useLegacyInstructions {
-		return nil, errors.New("only legacy transactions are supported")
-	}
-
 	// todo: struct this
 	reqBody := fmt.Sprintf(
 		`{"quoteResponse": %s, "userPublicKey": "%s", "destinationTokenAccount": "%s", "prioritizationFeeLamports": "auto", "asLegacyTransaction": %v}`,
@@ -206,6 +207,8 @@ func (c *Client) GetSwapInstructions(
 		}
 	}
 
+	res.AddressLookupTables = jsonBody.AddressLookupTableAddresses
+
 	return &res, nil
 }
 
@@ -258,9 +261,10 @@ type jsonInstruction struct {
 }
 
 type jsonSwapInstructions struct {
-	TokenLedgerInstruction    *jsonInstruction   `json:"tokenLedgerInstruction"`
-	ComputeBudgetInstructions []*jsonInstruction `json:"computeBudgetInstructions"`
-	SetupInstructions         []*jsonInstruction `json:"setupInstructions"`
-	SwapInstruction           *jsonInstruction   `json:"swapInstruction"`
-	CleanupInstruction        *jsonInstruction   `json:"cleanupInstruction"`
+	TokenLedgerInstruction      *jsonInstruction   `json:"tokenLedgerInstruction"`
+	ComputeBudgetInstructions   []*jsonInstruction `json:"computeBudgetInstructions"`
+	SetupInstructions           []*jsonInstruction `json:"setupInstructions"`
+	SwapInstruction             *jsonInstruction   `json:"swapInstruction"`
+	CleanupInstruction          *jsonInstruction   `json:"cleanupInstruction"`
+	AddressLookupTableAddresses []string           `json:"addressLookupTableAddresses"`
 }

pkg/solana/address_lookup_table.go

@@ -0,0 +1,25 @@
+package solana
+
+import (
+	"bytes"
+	"crypto/ed25519"
+)
+
+type AddressLookupTable struct {
+	PublicKey ed25519.PublicKey
+	Addresses []ed25519.PublicKey
+}
+
+type SortableAddressLookupTables []AddressLookupTable
+
+func (s SortableAddressLookupTables) Len() int {
+	return len(s)
+}
+
+func (s SortableAddressLookupTables) Less(i int, j int) bool {
+	return bytes.Compare(s[i].PublicKey, s[j].PublicKey) < 0
+}
+
+func (s SortableAddressLookupTables) Swap(i int, j int) {
+	s[i], s[j] = s[j], s[i]
+}

pkg/solana/encoding.go

@@ -49,6 +49,17 @@ func (t *Transaction) Unmarshal(b []byte) error {
 }
 
 func (m Message) Marshal() []byte {
+	switch m.version {
+	case MessageVersionLegacy:
+		return m.marshalLegacy()
+	case MessageVersion0:
+		return m.marshalV0()
+	default:
+		panic("unsupported message version")
+	}
+}
+
+func (m Message) marshalLegacy() []byte {
 	b := bytes.NewBuffer(nil)
 
 	// Header
@@ -82,7 +93,60 @@ func (m Message) Marshal() []byte {
 	return b.Bytes()
 }
 
+func (m Message) marshalV0() []byte {
+	b := bytes.NewBuffer(nil)
+
+	// Version Number
+	_ = b.WriteByte(byte(m.version + 127))
+
+	// Header
+	_ = b.WriteByte(m.Header.NumSignatures)
+	_ = b.WriteByte(m.Header.NumReadonlySigned)
+	_ = b.WriteByte(m.Header.NumReadOnly)
+
+	// Accounts
+	_, _ = shortvec.EncodeLen(b, len(m.Accounts))
+	for _, a := range m.Accounts {
+		_, _ = b.Write(a)
+	}
+
+	// Recent Blockhash
+	_, _ = b.Write(m.RecentBlockhash[:])
+
+	// Instructions
+	_, _ = shortvec.EncodeLen(b, len(m.Instructions))
+	for _, i := range m.Instructions {
+		_ = b.WriteByte(i.ProgramIndex)
+
+		// Accounts
+		_, _ = shortvec.EncodeLen(b, len(i.Accounts))
+		_, _ = b.Write(i.Accounts)
+
+		// Data
+		_, _ = shortvec.EncodeLen(b, len(i.Data))
+		_, _ = b.Write(i.Data)
+	}
+
+	_, _ = shortvec.EncodeLen(b, len(m.AddressTableLookups))
+	for _, addressTableLookup := range m.AddressTableLookups {
+		_, _ = b.Write(addressTableLookup.PublicKey)
+
+		_, _ = shortvec.EncodeLen(b, len(addressTableLookup.WritableIndexes))
+		_, _ = b.Write(addressTableLookup.WritableIndexes)
+
+		_, _ = shortvec.EncodeLen(b, len(addressTableLookup.ReadonlyIndexes))
+		_, _ = b.Write(addressTableLookup.ReadonlyIndexes)
+	}
+
+	return b.Bytes()
+}
+
 func (m *Message) Unmarshal(b []byte) (err error) {
+	// todo: double check this is correct
+	if b[0] > 127 {
+		return errors.New("versioned messages not supported")
+	}
+
 	buf := bytes.NewBuffer(b)
 
 	// Header

pkg/solana/transaction.go

@@ -20,17 +20,32 @@ const (
 type Signature [ed25519.SignatureSize]byte
 type Blockhash [sha256.Size]byte
 
+type MessageVersion uint8
+
+const (
+	MessageVersionLegacy MessageVersion = iota
+	MessageVersion0
+)
+
 type Header struct {
 	NumSignatures     byte
 	NumReadonlySigned byte
 	NumReadOnly       byte
 }
 
+type MessageAddressTableLookup struct {
+	PublicKey       ed25519.PublicKey
+	WritableIndexes []byte
+	ReadonlyIndexes []byte
+}
+
 type Message struct {
-	Header          Header
-	Accounts        []ed25519.PublicKey
-	RecentBlockhash Blockhash
-	Instructions    []CompiledInstruction
+	version             MessageVersion
+	Header              Header
+	Accounts            []ed25519.PublicKey
+	RecentBlockhash     Blockhash
+	Instructions        []CompiledInstruction
+	AddressTableLookups []MessageAddressTableLookup
 }
 
 type Transaction struct {
@@ -39,6 +54,15 @@ type Transaction struct {
 }
 
 func NewTransaction(payer ed25519.PublicKey, instructions ...Instruction) Transaction {
+	return newTransaction(payer, nil, instructions)
+}
+
+func NewVersionedTransaction(payer ed25519.PublicKey, addressLookupTables []AddressLookupTable, instructions []Instruction) Transaction {
+	return newTransaction(payer, addressLookupTables, instructions)
+}
+
+// todo: consolidate to new constructor
+func newTransaction(payer ed25519.PublicKey, addressLookupTables []AddressLookupTable, instructions []Instruction) Transaction {
 	accounts := []AccountMeta{
 		{
 			PublicKey:  payer,
@@ -65,8 +89,45 @@ func NewTransaction(payer ed25519.PublicKey, instructions ...Instruction) Transa
 	accounts = filterUnique(accounts)
 	sort.Sort(SortableAccountMeta(accounts))
 
+	// Sort address tables to guarantee consistent marshalling
+	sortedAddressLookupTables := make([]AddressLookupTable, len(addressLookupTables))
+	copy(sortedAddressLookupTables, addressLookupTables)
+	sort.Sort(SortableAddressLookupTables(sortedAddressLookupTables))
+
+	writableAddressTableIndexes := make([][]byte, len(sortedAddressLookupTables))
+	readonlyAddressTableIndexes := make([][]byte, len(sortedAddressLookupTables))
+
 	var m Message
 	for _, account := range accounts {
+		// If the account is eligible for dynamic loading, then pull its index
+		// from the first address table where it's defined.
+		var isDynamicallyLoaded bool
+		if !account.isPayer && !account.IsSigner && !account.isProgram {
+			for i, addressLookupTable := range sortedAddressLookupTables {
+				for j, address := range addressLookupTable.Addresses {
+					if bytes.Equal(address, account.PublicKey) {
+						isDynamicallyLoaded = true
+
+						if account.IsWritable {
+							writableAddressTableIndexes[i] = append(writableAddressTableIndexes[i], byte(j))
+						} else {
+							readonlyAddressTableIndexes[i] = append(readonlyAddressTableIndexes[i], byte(j))
+						}
+
+						break
+					}
+				}
+
+				if isDynamicallyLoaded {
+					break
+				}
+			}
+		}
+		if isDynamicallyLoaded {
+			continue
+		}
+
+		// Otherwise, the account is defined statically
 		m.Accounts = append(m.Accounts, account.PublicKey)
 
 		if account.IsSigner {
@@ -80,21 +141,58 @@ func NewTransaction(payer ed25519.PublicKey, instructions ...Instruction) Transa
 		}
 	}
 
+	// Consolidate static and dynamically loaded accounts into an ordered list,
+	// which is used for index references encoded in the message
+	dynamicWritableAccounts := make([]ed25519.PublicKey, 0)
+	dynamicReadonlyAccount := make([]ed25519.PublicKey, 0)
+	for i, writableAddressTableIndexes := range writableAddressTableIndexes {
+		for _, index := range writableAddressTableIndexes {
+			writableAccount := sortedAddressLookupTables[i].Addresses[index]
+			dynamicWritableAccounts = append(dynamicWritableAccounts, writableAccount)
+		}
+	}
+	for i, readonlyAddressTableIndexes := range readonlyAddressTableIndexes {
+		for _, index := range readonlyAddressTableIndexes {
+			readonlyAccount := sortedAddressLookupTables[i].Addresses[index]
+			dynamicReadonlyAccount = append(dynamicReadonlyAccount, readonlyAccount)
+		}
+	}
+	var allAccounts []ed25519.PublicKey
+	allAccounts = append(allAccounts, m.Accounts...)
+	allAccounts = append(allAccounts, dynamicWritableAccounts...)
+	allAccounts = append(allAccounts, dynamicReadonlyAccount...)
+
 	// Generate the compiled instruction, which uses indices instead
 	// of raw account keys.
 	for _, i := range instructions {
 		c := CompiledInstruction{
-			ProgramIndex: byte(indexOf(m.Accounts, i.Program)),
+			ProgramIndex: byte(indexOf(allAccounts, i.Program)),
 			Data:         i.Data,
 		}
 
 		for _, a := range i.Accounts {
-			c.Accounts = append(c.Accounts, byte(indexOf(m.Accounts, a.PublicKey)))
+			c.Accounts = append(c.Accounts, byte(indexOf(allAccounts, a.PublicKey)))
 		}
 
 		m.Instructions = append(m.Instructions, c)
 	}
 
+	// Generate the compiled message address table lookups
+	for i, addressLookupTable := range sortedAddressLookupTables {
+		if len(writableAddressTableIndexes[i]) == 0 && len(readonlyAddressTableIndexes[i]) == 0 {
+			continue
+		}
+
+		m.AddressTableLookups = append(m.AddressTableLookups, MessageAddressTableLookup{
+			PublicKey:       addressLookupTable.PublicKey,
+			WritableIndexes: writableAddressTableIndexes[i],
+			ReadonlyIndexes: readonlyAddressTableIndexes[i],
+		})
+	}
+	if len(m.AddressTableLookups) > 0 {
+		m.version = MessageVersion0
+	}
+
 	for i := range m.Accounts {
 		if len(m.Accounts[i]) == 0 {
 			m.Accounts[i] = make([]byte, ed25519.PublicKeySize)
@@ -118,11 +216,12 @@ func (t *Transaction) String() string {
 		sb.WriteString(fmt.Sprintf("  %d: %s\n", i, base58.Encode(s[:])))
 	}
 	sb.WriteString("Message:\n")
+	sb.WriteString(fmt.Sprintf("  Version: %s\n", t.Message.version.String()))
 	sb.WriteString("  Header:\n")
 	sb.WriteString(fmt.Sprintf("    NumSignatures: %d\n", t.Message.Header.NumSignatures))
 	sb.WriteString(fmt.Sprintf("    NumReadOnly: %d\n", t.Message.Header.NumReadOnly))
 	sb.WriteString(fmt.Sprintf("    NumReadOnlySigned: %d\n", t.Message.Header.NumReadonlySigned))
-	sb.WriteString("  Accounts:\n")
+	sb.WriteString("  Static Accounts:\n")
 	for i, a := range t.Message.Accounts {
 		sb.WriteString(fmt.Sprintf("    %d: %s\n", i, base58.Encode(a)))
 	}
@@ -133,7 +232,14 @@ func (t *Transaction) String() string {
 		sb.WriteString(fmt.Sprintf("      Accounts: %v\n", t.Message.Instructions[i].Accounts))
 		sb.WriteString(fmt.Sprintf("      Data: %v\n", t.Message.Instructions[i].Data))
 	}
-
+	if len(t.Message.AddressTableLookups) > 0 {
+		sb.WriteString("  Address Table Lookups:\n")
+		for i := range t.Message.AddressTableLookups {
+			sb.WriteString(fmt.Sprintf("    %s:\n", base58.Encode(t.Message.AddressTableLookups[i].PublicKey)))
+			sb.WriteString(fmt.Sprintf("      Writable Indexes: %v\n", t.Message.AddressTableLookups[i].WritableIndexes))
+			sb.WriteString(fmt.Sprintf("      Readonly Indexes: %v\n", t.Message.AddressTableLookups[i].ReadonlyIndexes))
+		}
+	}
 	return sb.String()
 }
 
@@ -198,3 +304,13 @@ func indexOf(slice []ed25519.PublicKey, item ed25519.PublicKey) int {
 
 	return -1
 }
+
+func (v MessageVersion) String() string {
+	switch v {
+	case MessageVersionLegacy:
+		return "legacy"
+	case MessageVersion0:
+		return "v0"
+	}
+	return "unknown"
+}