Add decode for icrc accounts.

aviate-labs · Jun 8, 2023 · 0641394 · 0641394
1 parent f6e7ebe
commit 0641394
Show file tree

Hide file tree

Showing 5 changed files with 157 additions and 17 deletions.
diff --git a/principal/accountid.go b/principal/accountid.go
@@ -4,6 +4,7 @@ import (
 	"crypto/sha256"
 	"encoding/binary"
 	"encoding/hex"
+	"fmt"
 	"hash/crc32"
 )
 
@@ -19,8 +20,8 @@ type (
 	SubAccount [32]byte
 )
 
-// FromPrincipal returns the account identifier corresponding with the given sub-account.
-func FromPrincipal(p Principal, subAccount [32]byte) AccountIdentifier {
+// NewAccountID returns the account identifier corresponding with the given sub-account.
+func NewAccountID(p Principal, subAccount [32]byte) AccountIdentifier {
 	h := sha256.New224()
 	h.Write([]byte("\x0Aaccount-id"))
 	h.Write(p.Raw)
@@ -39,7 +40,30 @@ func (id AccountIdentifier) Bytes() []byte {
 	return append(crc, id[:]...)
 }
 
+// DecodeAccountID decodes the given string into an account identifier.
+func DecodeAccountID(s string) (AccountIdentifier, error) {
+	b, err := hex.DecodeString(s)
+	if err != nil {
+		return AccountIdentifier{}, err
+	}
+	if len(b) != 32 {
+		return AccountIdentifier{}, fmt.Errorf("invalid length: %d", len(b))
+	}
+	crc := binary.BigEndian.Uint32(b[:4])
+	if crc != crc32.ChecksumIEEE(b[4:]) {
+		return AccountIdentifier{}, fmt.Errorf("invalid checksum: %d", crc)
+	}
+	var id AccountIdentifier
+	copy(id[:], b[4:])
+	return id, nil
+}
+
+// Encode returns the hexadecimal representation of the account identifier.
+func (id AccountIdentifier) Encode() string {
+	return hex.EncodeToString(id.Bytes())
+}
+
 // String returns the hexadecimal representation of the account identifier.
 func (id AccountIdentifier) String() string {
-	return hex.EncodeToString(id.Bytes())
+	return id.Encode()
 }
diff --git a/principal/accountid_test.go b/principal/accountid_test.go
@@ -3,10 +3,24 @@ package principal_test
 import (
 	"fmt"
 	"github.com/aviate-labs/agent-go/principal"
+	"testing"
 )
 
-func ExampleFromPrincipal() {
-	fmt.Println(principal.FromPrincipal(principal.Principal{}, principal.DefaultSubAccount).String())
+func ExampleNewAccountID() {
+	fmt.Println(principal.NewAccountID(principal.Principal{}, principal.DefaultSubAccount).String())
 	// Output:
 	// 2d0e897f7e862d2b57d9bc9ea5c65f9a24ac6c074575f47898314b8d6cb0929d
 }
+
+func TestAccountIdentifier(t *testing.T) {
+	for i := 0; i < 100; i++ {
+		a := principal.NewAccountID(principal.Principal{}, [32]byte{byte(i)})
+		accountID, err := principal.DecodeAccountID(a.String())
+		if err != nil {
+			t.Error(err)
+		}
+		if accountID != a {
+			t.Errorf("expected %v, got %v", a, accountID)
+		}
+	}
+}
diff --git a/principal/icrc/account.go b/principal/icrc/account.go
@@ -4,11 +4,14 @@ import (
 	"encoding/base32"
 	"encoding/binary"
 	"encoding/hex"
+	"fmt"
 	"github.com/aviate-labs/agent-go/principal"
 	"hash/crc32"
 	"strings"
 )
 
+var encoding = base32.StdEncoding.WithPadding(base32.NoPadding)
+
 func trimLeadingZeros(str string) string {
 	for str[0] == '0' {
 		str = str[1:]
@@ -21,7 +24,60 @@ type Account struct {
 	SubAccount *[32]byte
 }
 
-func (a Account) String() string {
+func Decode(s string) (Account, error) {
+	a := strings.Split(s, ".")
+	if len(a) == 1 {
+		owner, err := principal.Decode(s)
+		if err != nil {
+			return Account{}, err
+		}
+		return Account{
+			Owner:      owner,
+			SubAccount: nil,
+		}, nil
+	}
+	if len(a) != 2 {
+		return Account{}, fmt.Errorf("invalid account identifier: %s", s)
+	}
+	p := strings.Split(a[0], "-")
+	b32crc := strings.ToUpper(p[len(p)-1])
+	owner, err := principal.Decode(strings.Join(p[:len(p)-1], "-"))
+	if err != nil {
+		return Account{}, err
+	}
+	if len(a[1]) == 0 || a[1][0] == '0' {
+		return Account{}, fmt.Errorf("invalid sub account: %s", a[1])
+	}
+	if len(a[1])%2 == 1 {
+		// Add leading zero if necessary.
+		a[1] = "0" + a[1]
+	}
+	subAccount, err := hex.DecodeString(a[1])
+	if err != nil {
+		return Account{}, err
+	}
+	for len(subAccount) < 32 {
+		subAccount = append([]byte{0}, subAccount...)
+	}
+	cs, err := encoding.DecodeString(b32crc)
+	if err != nil {
+		return Account{}, err
+	}
+	if len(cs) != 4 {
+		return Account{}, fmt.Errorf("invalid checksum size: %d", len(cs))
+	}
+	if crc32.ChecksumIEEE(append(owner.Raw, subAccount...)) != binary.BigEndian.Uint32(cs) {
+		return Account{}, fmt.Errorf("invalid checksum: %s", string(cs))
+	}
+	var subAccount32 [32]byte
+	copy(subAccount32[:], subAccount)
+	return Account{
+		Owner:      owner,
+		SubAccount: &subAccount32,
+	}, nil
+}
+
+func (a Account) Encode() string {
 	if a.SubAccount == nil {
 		return a.Owner.String()
 	}
@@ -30,13 +86,10 @@ func (a Account) String() string {
 	}
 	cs := make([]byte, 4)
 	binary.BigEndian.PutUint32(cs, crc32.ChecksumIEEE(append(a.Owner.Raw, a.SubAccount[:]...)))
-	b32cs := strings.ToLower(removePadding(base32.StdEncoding.EncodeToString(cs)))
+	b32cs := strings.ToLower(encoding.EncodeToString(cs))
 	return a.Owner.String() + "-" + b32cs + "." + trimLeadingZeros(hex.EncodeToString(a.SubAccount[:]))
 }
 
-func removePadding(str string) string {
-	for strings.HasSuffix(str, "=") {
-		str = str[:len(str)-1]
-	}
-	return str
+func (a Account) String() string {
+	return a.Encode()
 }
diff --git a/principal/icrc/account_test.go b/principal/icrc/account_test.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"github.com/aviate-labs/agent-go/principal"
 	"github.com/aviate-labs/agent-go/principal/icrc"
+	"testing"
 )
 
 func ExampleAccount() {
@@ -19,3 +20,45 @@ func ExampleAccount() {
 	// Output:
 	// k2t6j-2nvnp-4zjm3-25dtz-6xhaa-c7boj-5gayf-oj3xs-i43lp-teztq-6ae-dfxgiyy.102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20
 }
+
+func TestVectors(t *testing.T) {
+	for _, test := range []struct {
+		account string
+		err     bool
+	}{
+		{account: "k2t6j-2nvnp-4zjm3-25dtz-6xhaa-c7boj-5gayf-oj3xs-i43lp-teztq-6ae"},
+		{account: "k2t6j-2nvnp-4zjm3-25dtz-6xhaa-c7boj-5gayf-oj3xs-i43lp-teztq-6ae-q6bn32y.", err: true},
+		{account: "k2t6j2nvnp4zjm3-25dtz6xhaac7boj5gayfoj3xs-i43lp-teztq-6ae", err: true},
+		{account: "k2t6j-2nvnp-4zjm3-25dtz-6xhaa-c7boj-5gayf-oj3xs-i43lp-teztq-6ae-6cc627i.1"},
+		{account: "k2t6j-2nvnp-4zjm3-25dtz-6xhaa-c7boj-5gayf-oj3xs-i43lp-teztq-6ae-6cc627i.01", err: true},
+		{account: "k2t6j-2nvnp-4zjm3-25dtz-6xhaa-c7boj-5gayf-oj3xs-i43lp-teztq-6ae.1", err: true},
+		{account: "k2t6j-2nvnp-4zjm3-25dtz-6xhaa-c7boj-5gayf-oj3xs-i43lp-teztq-6ae-dfxgiyy.102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20"},
+	} {
+		a, err := icrc.Decode(test.account)
+		if err != nil {
+			if !test.err {
+				t.Error(err)
+			}
+			continue
+		}
+		if a.String() != test.account {
+			t.Errorf("expected %s, got %s", test.account, a.String())
+		}
+	}
+}
+
+func TestAccountIdentifier(t *testing.T) {
+	for i := 0; i < 100; i++ {
+		a := icrc.Account{
+			Owner:      principal.Principal{},
+			SubAccount: &[32]byte{byte(i)},
+		}
+		accountID, err := icrc.Decode(a.String())
+		if err != nil {
+			t.Error(err)
+		}
+		if accountID.String() != a.String() {
+			t.Errorf("expected %s, got %s", a.String(), accountID.String())
+		}
+	}
+}
diff --git a/principal/principal.go b/principal/principal.go
@@ -22,9 +22,16 @@ type Principal struct {
 
 // Decode converts a textual representation into a principal.
 func Decode(s string) (Principal, error) {
-	s = strings.ReplaceAll(s, "-", "")
-	s = strings.ToUpper(s)
-	b32, err := encoding.DecodeString(s)
+	p := strings.Split(s, "-")
+	for i, c := range p {
+		if len(c) > 5 {
+			return Principal{}, fmt.Errorf("invalid length: %s", c)
+		}
+		if i != len(p)-1 && len(c) < 5 {
+			return Principal{}, fmt.Errorf("invalid length: %s", c)
+		}
+	}
+	b32, err := encoding.DecodeString(strings.ToUpper(strings.Join(p, "")))
 	if err != nil {
 		return Principal{}, err
 	}
@@ -49,8 +56,7 @@ func NewSelfAuthenticating(pub []byte) Principal {
 func (p Principal) Encode() string {
 	cs := make([]byte, 4)
 	binary.BigEndian.PutUint32(cs, crc32.ChecksumIEEE(p.Raw))
-	b32 := encoding.EncodeToString(append(cs, p.Raw...))
-	b32 = strings.ToLower(b32)
+	b32 := strings.ToLower(encoding.EncodeToString(append(cs, p.Raw...)))
 	var str string
 	for i, c := range b32 {
 		if i != 0 && i%5 == 0 {