From 6f37c3b5778a491bebce6745a92d9d07bc4f6ff7 Mon Sep 17 00:00:00 2001
From: Andrew Chiw <randomshinichi4869@gmail.com>
Date: Wed, 6 Nov 2019 15:59:01 +0100
Subject: [PATCH] refactor: restrict HD wallet functionality to only deriving
 aeternity addresses

---
 account/hdwallet.go      | 25 ++++++++++++++++++------
 account/hdwallet_test.go | 41 +++++++++++++++++++++++-----------------
 2 files changed, 43 insertions(+), 23 deletions(-)

diff --git a/account/hdwallet.go b/account/hdwallet.go
index 59f92ddc..6b72bcc8 100644
--- a/account/hdwallet.go
+++ b/account/hdwallet.go
@@ -2,6 +2,7 @@ package account
 
 import (
 	"bytes"
+	"fmt"
 
 	"github.com/tyler-smith/go-bip39"
 	"golang.org/x/crypto/nacl/sign"
@@ -27,27 +28,39 @@ func BIP32KeyToAeKey(key *Key) (acc *Account, err error) {
 }
 
 // DerivePathFromSeed derives a BIP32 Key given a seed (usually derived from a
-// mnemonic) and a path. Due to ed25519, only hardened path nodes are supported.
-// Hardened nodes are denoted with apostrophes ', e.g. "m/44'/457'/0'/0'/0'".
-func DerivePathFromSeed(masterSeed []byte, path string) (key *Key, err error) {
+// mnemonic) and an account index and address index, as standardized by AEX-10.
+func DerivePathFromSeed(masterSeed []byte, accountIndex, addressIndex uint32) (key *Key, err error) {
+	path := buildPath(accountIndex, addressIndex)
 	mK, err := NewMasterKey(masterSeed)
 	if err != nil {
 		return
 	}
+
+	return derivePath(mK, path)
+}
+
+// buildPath ensures that only aeternity BIP32 addresses can be calculated. Due
+// to ed25519, only hardened path nodes are supported. Hardened nodes are
+// denoted with apostrophes ', e.g. "m/44'/457'/0'/0'/0'".
+func buildPath(accountIndex, addressIndex uint32) string {
+	return fmt.Sprintf("m/44'/457'/%d'/0'/%d'", accountIndex, addressIndex)
+}
+
+func derivePath(key *Key, path string) (childKey *Key, err error) {
 	parsedPath, err := ParsePath(path)
 	if err != nil {
 		return
 	}
 
-	key = mK
+	childKey = key
 	for _, p := range parsedPath.Elements {
 		if p.Master {
 			continue
 		}
-		key, err = key.NewChildKey(p.ChildNumber)
+		childKey, err = childKey.NewChildKey(p.ChildNumber)
 		if err != nil {
 			return
 		}
 	}
-	return key, nil
+	return
 }
diff --git a/account/hdwallet_test.go b/account/hdwallet_test.go
index c845d3b6..8b9a0e52 100644
--- a/account/hdwallet_test.go
+++ b/account/hdwallet_test.go
@@ -8,31 +8,34 @@ import (
 )
 
 func TestDerivePathFromSeedAeAccount(t *testing.T) {
+	type args struct {
+		accountIndex uint32
+		addressIndex uint32
+	}
 	tests := []struct {
 		name        string
 		seed        string
-		path        string
+		args        args
 		wantAddress string
 		wantErr     bool
 	}{
 		{
-			name:        "Standard Seed, m/0'",
-			seed:        "60812c7c93d6f9cb346bbcf799957b6ec776aea84b01bdd9f9b7916522cc52c6ea5d07960b68668cd37b0a77f6c4fe283f146bd916153c426df126a8b8707b39",
-			path:        "m/0'",
-			wantAddress: "ak_2o3zmfG3hMFu4oveTs4VcsmdimnsUBX7sEp3LwXhj9kutRm8mN",
-			wantErr:     false,
-		},
-		{
-			name:        "Standard Seed, m/44'/457'/0'/0'/0'",
-			seed:        "60812c7c93d6f9cb346bbcf799957b6ec776aea84b01bdd9f9b7916522cc52c6ea5d07960b68668cd37b0a77f6c4fe283f146bd916153c426df126a8b8707b39",
-			path:        "m/44'/457'/0'/0'/0'",
+			name: "Standard Seed, m/44'/457'/0'/0'/0'",
+			seed: "60812c7c93d6f9cb346bbcf799957b6ec776aea84b01bdd9f9b7916522cc52c6ea5d07960b68668cd37b0a77f6c4fe283f146bd916153c426df126a8b8707b39",
+			args: args{
+				accountIndex: 0,
+				addressIndex: 0,
+			},
 			wantAddress: "ak_2Z74Jhbo3xqF47k2h6NoUpr5gTfc9EQFX7wPH2Vf7Q5PCVcZSW",
 			wantErr:     false,
 		},
 		{
-			name:        "Standard Seed, m/44'/457'/0'/0'/3'",
-			seed:        "60812c7c93d6f9cb346bbcf799957b6ec776aea84b01bdd9f9b7916522cc52c6ea5d07960b68668cd37b0a77f6c4fe283f146bd916153c426df126a8b8707b39",
-			path:        "m/44'/457'/0'/0'/3'",
+			name: "Standard Seed, m/44'/457'/0'/0'/3'",
+			seed: "60812c7c93d6f9cb346bbcf799957b6ec776aea84b01bdd9f9b7916522cc52c6ea5d07960b68668cd37b0a77f6c4fe283f146bd916153c426df126a8b8707b39",
+			args: args{
+				accountIndex: 0,
+				addressIndex: 3,
+			},
 			wantAddress: "ak_2wPpjbxDhdn8PURqLPsunqBTWYSbe9iac1gjfJcQVzY4aZYEzq",
 			wantErr:     false,
 		},
@@ -42,7 +45,7 @@ func TestDerivePathFromSeedAeAccount(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		key, err := DerivePathFromSeed(seedBytes, tt.path)
+		key, err := DerivePathFromSeed(seedBytes, tt.args.accountIndex, tt.args.addressIndex)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -57,7 +60,7 @@ func TestDerivePathFromSeedAeAccount(t *testing.T) {
 	}
 }
 
-func TestDerivePathFromSeed_SLIP0010_TestVectors(t *testing.T) {
+func TestDerivePath_SLIP0010_TestVectors(t *testing.T) {
 	// BIP32 adds an extra 0 byte of padding in the beginning of the public key.
 	// We do not use the padded public key. Instead we just feed the private key
 	// into the official ed25519 scheme and use its public key. Nevertheless it
@@ -185,7 +188,11 @@ func TestDerivePathFromSeed_SLIP0010_TestVectors(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		key, err := DerivePathFromSeed(seedBytes, tt.path)
+		masterKey, err := NewMasterKey(seedBytes)
+		if err != nil {
+			t.Fatal(err)
+		}
+		key, err := derivePath(masterKey, tt.path)
 		if err != nil {
 			t.Fatal(err)
 		}