From 4a3ecddb010b88cee5d5317abd23845ed4f67a00 Mon Sep 17 00:00:00 2001
From: Richard Patel <richard@blockdaemon.com>
Date: Tue, 22 Feb 2022 10:55:15 +0100
Subject: [PATCH] allow PublicKey to be used as flag

---
 keys.go      | 16 ++++++++++------
 keys_test.go | 15 +++++++++++++++
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/keys.go b/keys.go
index a4cd237d..1ba32795 100644
--- a/keys.go
+++ b/keys.go
@@ -144,12 +144,8 @@ func (p PublicKey) MarshalText() ([]byte, error) {
 	return []byte(base58.Encode(p[:])), nil
 }
 
-func (p *PublicKey) UnmarshalText(data []byte) (err error) {
-	*p, err = PublicKeyFromBase58(string(data))
-	if err != nil {
-		return fmt.Errorf("invalid public key %q: %w", data, err)
-	}
-	return
+func (p *PublicKey) UnmarshalText(data []byte) error {
+	return p.Set(string(data))
 }
 
 func (p PublicKey) MarshalJSON() ([]byte, error) {
@@ -195,6 +191,14 @@ func (p PublicKey) IsZero() bool {
 	return p == zeroPublicKey
 }
 
+func (p *PublicKey) Set(s string) (err error) {
+	*p, err = PublicKeyFromBase58(s)
+	if err != nil {
+		return fmt.Errorf("invalid public key %s: %w", s, err)
+	}
+	return
+}
+
 func (p PublicKey) String() string {
 	return base58.Encode(p[:])
 }
diff --git a/keys_test.go b/keys_test.go
index 53f407dc..a37ff84c 100644
--- a/keys_test.go
+++ b/keys_test.go
@@ -21,6 +21,7 @@ import (
 	"encoding/binary"
 	"encoding/hex"
 	"errors"
+	"flag"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -159,6 +160,20 @@ func TestPublicKey_MarshalText(t *testing.T) {
 	)
 }
 
+func TestPublicKey_Flag(t *testing.T) {
+	flagSet := flag.NewFlagSet("", flag.ContinueOnError)
+	var key PublicKey
+	flagSet.Var(&key, "account", "Public key")
+	err := flagSet.Parse([]string{"--account", "7cVfgArCheMR6Cs4t6vz5rfnqd56vZq4ndaBrY5xkxXy"})
+	require.NoError(t, err)
+	assert.Equal(t, PublicKey{
+		0x62, 0x3d, 0xdd, 0x11, 0x7e, 0x7c, 0xc5, 0x62,
+		0xf6, 0x63, 0x15, 0x05, 0x25, 0x8c, 0xd1, 0xdc,
+		0xee, 0x81, 0x94, 0x9f, 0x8a, 0xfd, 0x1e, 0xa2,
+		0x94, 0xdc, 0x47, 0xbe, 0x6e, 0xcf, 0xf3, 0xa8,
+	}, key)
+}
+
 func TestPublicKeySlice(t *testing.T) {
 	{
 		slice := make(PublicKeySlice, 0)