From 2b76a3015c852c76773ae6c896b1b678ee64b4ec Mon Sep 17 00:00:00 2001
From: Anirudha Bose <anirudha.bose@alumni.cern>
Date: Fri, 28 Aug 2020 01:14:04 +0200
Subject: [PATCH] hdkeychain: add CloneWithVersion to set custom HD version
 bytes

This adds a new method to the ExtendedKey type that allows cloning the
extended key with custom HD version bytes. It does not mutate the
original extended key on which the method is called.

Added some tests to demonstrate the utility of this method, i.e.,
conversion between standard and SLIP-0132 extended keys.
---
 hdkeychain/extendedkey.go      | 30 +++++++++++++++
 hdkeychain/extendedkey_test.go | 68 ++++++++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+)

diff --git a/hdkeychain/extendedkey.go b/hdkeychain/extendedkey.go
index 2b59f04c0..3fe51c17d 100644
--- a/hdkeychain/extendedkey.go
+++ b/hdkeychain/extendedkey.go
@@ -378,6 +378,36 @@ func (k *ExtendedKey) Neuter() (*ExtendedKey, error) {
 		k.depth, k.childNum, false), nil
 }
 
+// CloneWithVersion returns a new extended key cloned from this extended key,
+// but using the provided HD version bytes. The version must be a private HD
+// key ID for an extended private key, and a public HD key ID for an extended
+// public key.
+//
+// This method creates a new copy and therefore does not mutate the original
+// extended key instance.
+//
+// Unlike Neuter(), this does NOT convert an extended private key to an
+// extended public key. It is particularly useful for converting between
+// standard BIP0032 extended keys (serializable to xprv/xpub) and keys based
+// on the SLIP132 standard (serializable to yprv/ypub, zprv/zpub, etc.).
+//
+// References:
+//   [SLIP132]: SLIP-0132 - Registered HD version bytes for BIP-0032
+//   https://github.com/satoshilabs/slips/blob/master/slip-0132.md
+func (k *ExtendedKey) CloneWithVersion(version []byte) (*ExtendedKey, error) {
+	if len(version) != 4 {
+		// TODO: The semantically correct error to return here is
+		//  ErrInvalidHDKeyID (introduced in btcsuite/btcd#1617). Update the
+		//  error type once available in a stable btcd / chaincfg release.
+		return nil, chaincfg.ErrUnknownHDKeyID
+	}
+
+	// Initialize a new extended key instance with the same fields as the
+	// current extended private/public key and the provided HD version bytes.
+	return NewExtendedKey(version, k.key, k.chainCode, k.parentFP,
+		k.depth, k.childNum, k.isPrivate), nil
+}
+
 // ECPubKey converts the extended key to a btcec public key and returns it.
 func (k *ExtendedKey) ECPubKey() (*btcec.PublicKey, error) {
 	return btcec.ParsePubKey(k.pubKeyBytes(), btcec.S256())
diff --git a/hdkeychain/extendedkey_test.go b/hdkeychain/extendedkey_test.go
index 0fe49f213..868b840bf 100644
--- a/hdkeychain/extendedkey_test.go
+++ b/hdkeychain/extendedkey_test.go
@@ -1088,3 +1088,71 @@ func TestMaximumDepth(t *testing.T) {
 		t.Fatal("Child: deriving 256th key should not succeed")
 	}
 }
+
+// TestCloneWithVersion ensures proper conversion between standard and SLIP132
+// extended keys.
+//
+// The following tool was used for generating the tests:
+//   https://jlopp.github.io/xpub-converter
+func TestCloneWithVersion(t *testing.T) {
+	tests := []struct {
+		name    string
+		key     string
+		version []byte
+		want    string
+		wantErr error
+	}{
+		{
+			name:    "test xpub to zpub",
+			key:     "xpub661MyMwAqRbcFtXgS5sYJABqqG9YLmC4Q1Rdap9gSE8NqtwybGhePY2gZ29ESFjqJoCu1Rupje8YtGqsefD265TMg7usUDFdp6W1EGMcet8",
+			version: []byte{0x04, 0xb2, 0x47, 0x46},
+			want:    "zpub6jftahH18ngZxUuv6oSniLNrBCSSE1B4EEU59bwTCEt8x6aS6b2mdfLxbS4QS53g85SWWP6wexqeer516433gYpZQoJie2tcMYdJ1SYYYAL",
+		},
+		{
+			name:    "test zpub to xpub",
+			key:     "zpub6jftahH18ngZxUuv6oSniLNrBCSSE1B4EEU59bwTCEt8x6aS6b2mdfLxbS4QS53g85SWWP6wexqeer516433gYpZQoJie2tcMYdJ1SYYYAL",
+			version: []byte{0x04, 0x88, 0xb2, 0x1e},
+			want:    "xpub661MyMwAqRbcFtXgS5sYJABqqG9YLmC4Q1Rdap9gSE8NqtwybGhePY2gZ29ESFjqJoCu1Rupje8YtGqsefD265TMg7usUDFdp6W1EGMcet8",
+		},
+		{
+			name:    "test xprv to zprv",
+			key:     "xprv9s21ZrQH143K3QTDL4LXw2F7HEK3wJUD2nW2nRk4stbPy6cq3jPPqjiChkVvvNKmPGJxWUtg6LnF5kejMRNNU3TGtRBeJgk33yuGBxrMPHi",
+			version: []byte{0x04, 0xb2, 0x43, 0x0c},
+			want:    "zprvAWgYBBk7JR8GjzqSzmunMCS7dAbwpYTCs1YUMDXqduMA5JFHZ3iX5s2UkAR6vBdcCYYa1S5o1fVLrKsrnpCQ4WpUd6aVUWP1bS2Yy5DoaKv",
+		},
+		{
+			name:    "test zprv to xprv",
+			key:     "zprvAWgYBBk7JR8GjzqSzmunMCS7dAbwpYTCs1YUMDXqduMA5JFHZ3iX5s2UkAR6vBdcCYYa1S5o1fVLrKsrnpCQ4WpUd6aVUWP1bS2Yy5DoaKv",
+			version: []byte{0x04, 0x88, 0xad, 0xe4},
+			want:    "xprv9s21ZrQH143K3QTDL4LXw2F7HEK3wJUD2nW2nRk4stbPy6cq3jPPqjiChkVvvNKmPGJxWUtg6LnF5kejMRNNU3TGtRBeJgk33yuGBxrMPHi",
+		},
+		{
+			name:    "test invalid key id",
+			key:     "zprvAWgYBBk7JR8GjzqSzmunMCS7dAbwpYTCs1YUMDXqduMA5JFHZ3iX5s2UkAR6vBdcCYYa1S5o1fVLrKsrnpCQ4WpUd6aVUWP1bS2Yy5DoaKv",
+			version: []byte{0x4B, 0x1D},
+			wantErr: chaincfg.ErrUnknownHDKeyID,
+		},
+	}
+
+	for i, test := range tests {
+		extKey, err := NewKeyFromString(test.key)
+		if err != nil {
+			panic(err) // This is never expected to fail.
+		}
+
+		got, err := extKey.CloneWithVersion(test.version)
+		if !reflect.DeepEqual(err, test.wantErr) {
+			t.Errorf("CloneWithVersion #%d (%s): unexpected error -- "+
+				"want %v, got %v", i, test.name, test.wantErr, err)
+			continue
+		}
+
+		if test.wantErr == nil {
+			if k := got.String(); k != test.want {
+				t.Errorf("CloneWithVersion #%d (%s): "+
+					"got %s, want %s", i, test.name, k, test.want)
+				continue
+			}
+		}
+	}
+}