From bdab8dfe81c92667a2e2bf9044e068398be50549 Mon Sep 17 00:00:00 2001
From: Anirudha Bose <anirudha.bose@alumni.cern>
Date: Mon, 24 Aug 2020 20:48:23 +0200
Subject: [PATCH] chaincfg: Add RegisterHDKeyID func to populate HD key ID
 pairs

Currently, the only way to register HD version bytes is by initializing
chaincfg.Params struct, and registering it during package init.
RegisterHDKeyID provides a way to populate custom HD version bytes,
without having to create new chaincfg.Params instances. This is useful
for library packages who want to use non-standard version bytes for
serializing extended keys, such as the ones documented in SLIP-0132.

This function is complementary to HDPrivateKeyToPublicKeyID, which is
used to lookup previously registered key IDs.
---
 chaincfg/params.go      | 33 ++++++++++++++++++++++++-
 chaincfg/params_test.go | 53 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 84 insertions(+), 2 deletions(-)

diff --git a/chaincfg/params.go b/chaincfg/params.go
index 54117b8713..7e4327984c 100644
--- a/chaincfg/params.go
+++ b/chaincfg/params.go
@@ -588,6 +588,10 @@ var (
 	// is intended to identify the network for a hierarchical deterministic
 	// private extended key is not registered.
 	ErrUnknownHDKeyID = errors.New("unknown hd private extended key bytes")
+
+	// ErrInvalidHDKeyID describes an error where the provided hierarchical
+	// deterministic version bytes, or hd key id, is malformed.
+	ErrInvalidHDKeyID = errors.New("invalid hd extended key version bytes")
 )
 
 var (
@@ -619,7 +623,11 @@ func Register(params *Params) error {
 	registeredNets[params.Net] = struct{}{}
 	pubKeyHashAddrIDs[params.PubKeyHashAddrID] = struct{}{}
 	scriptHashAddrIDs[params.ScriptHashAddrID] = struct{}{}
-	hdPrivToPubKeyIDs[params.HDPrivateKeyID] = params.HDPublicKeyID[:]
+
+	err := RegisterHDKeyID(params.HDPublicKeyID[:], params.HDPrivateKeyID[:])
+	if err != nil {
+		return err
+	}
 
 	// A valid Bech32 encoded segwit address always has as prefix the
 	// human-readable part for the given net followed by '1'.
@@ -666,6 +674,29 @@ func IsBech32SegwitPrefix(prefix string) bool {
 	return ok
 }
 
+// RegisterHDKeyID registers a public and private hierarchical deterministic
+// extended key ID pair.
+//
+// Non-standard HD version bytes, such as the ones documented in SLIP-0132,
+// should be registered using this method for library packages to lookup key
+// IDs (aka HD version bytes). When the provided key IDs are invalid, the
+// ErrInvalidHDKeyID error will be returned.
+//
+// Reference:
+//   SLIP-0132 : Registered HD version bytes for BIP-0032
+//   https://github.com/satoshilabs/slips/blob/master/slip-0132.md
+func RegisterHDKeyID(hdPublicKeyID []byte, hdPrivateKeyID []byte) error {
+	if len(hdPublicKeyID) != 4 || len(hdPrivateKeyID) != 4 {
+		return ErrInvalidHDKeyID
+	}
+
+	var keyID [4]byte
+	copy(keyID[:], hdPrivateKeyID)
+	hdPrivToPubKeyIDs[keyID] = hdPublicKeyID
+
+	return nil
+}
+
 // HDPrivateKeyToPublicKeyID accepts a private hierarchical deterministic
 // extended key id and returns the associated public key id.  When the provided
 // id is not registered, the ErrUnknownHDKeyID error will be returned.
diff --git a/chaincfg/params_test.go b/chaincfg/params_test.go
index 277a56bdd5..d9e1516812 100644
--- a/chaincfg/params_test.go
+++ b/chaincfg/params_test.go
@@ -4,7 +4,10 @@
 
 package chaincfg
 
-import "testing"
+import (
+	"bytes"
+	"testing"
+)
 
 // TestInvalidHashStr ensures the newShaHashFromStr function panics when used to
 // with an invalid hash string.
@@ -33,3 +36,51 @@ func TestMustRegisterPanic(t *testing.T) {
 	// Intentionally try to register duplicate params to force a panic.
 	mustRegister(&MainNetParams)
 }
+
+func TestRegisterHDKeyID(t *testing.T) {
+	t.Parallel()
+
+	// Ref: https://github.com/satoshilabs/slips/blob/master/slip-0132.md
+	hdKeyIDZprv := []byte{0x02, 0xaa, 0x7a, 0x99}
+	hdKeyIDZpub := []byte{0x02, 0xaa, 0x7e, 0xd3}
+
+	if err := RegisterHDKeyID(hdKeyIDZpub, hdKeyIDZprv); err != nil {
+		t.Fatalf("RegisterHDKeyID: expected no error, got %v", err)
+	}
+
+	got, err := HDPrivateKeyToPublicKeyID(hdKeyIDZprv)
+	if err != nil {
+		t.Fatalf("HDPrivateKeyToPublicKeyID: expected no error, got %v", err)
+	}
+
+	if !bytes.Equal(got, hdKeyIDZpub) {
+		t.Fatalf("HDPrivateKeyToPublicKeyID: expected result %v, got %v",
+			hdKeyIDZpub, got)
+	}
+}
+
+func TestInvalidHDKeyID(t *testing.T) {
+	t.Parallel()
+
+	prvValid := []byte{0x02, 0xaa, 0x7a, 0x99}
+	pubValid := []byte{0x02, 0xaa, 0x7e, 0xd3}
+	prvInvalid := []byte{0x00}
+	pubInvalid := []byte{0x00}
+
+	if err := RegisterHDKeyID(pubInvalid, prvValid); err != ErrInvalidHDKeyID {
+		t.Fatalf("RegisterHDKeyID: want err ErrInvalidHDKeyID, got %v", err)
+	}
+
+	if err := RegisterHDKeyID(pubValid, prvInvalid); err != ErrInvalidHDKeyID {
+		t.Fatalf("RegisterHDKeyID: want err ErrInvalidHDKeyID, got %v", err)
+	}
+
+	if err := RegisterHDKeyID(pubInvalid, prvInvalid); err != ErrInvalidHDKeyID {
+		t.Fatalf("RegisterHDKeyID: want err ErrInvalidHDKeyID, got %v", err)
+	}
+
+	// FIXME: The error type should be changed to ErrInvalidHDKeyID.
+	if _, err := HDPrivateKeyToPublicKeyID(prvInvalid); err != ErrUnknownHDKeyID {
+		t.Fatalf("HDPrivateKeyToPublicKeyID: want err ErrUnknownHDKeyID, got %v", err)
+	}
+}