From 3edeb0cd45a9a8a97dbafb901894a51c8f8af688 Mon Sep 17 00:00:00 2001
From: Ethan Frey <ethanfrey@users.noreply.github.com>
Date: Wed, 13 Sep 2017 11:53:14 +0200
Subject: [PATCH] Ledger caches pubkey, works with newer firmware

---
 nano/keys.go      | 39 +++++++++++++++++++++++++++++++++++----
 nano/keys_test.go | 10 ++++++----
 2 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/nano/keys.go b/nano/keys.go
index 55af915cb..7cac5eee8 100644
--- a/nano/keys.go
+++ b/nano/keys.go
@@ -4,6 +4,8 @@ import (
 	"bytes"
 	"encoding/hex"
 
+	"github.com/pkg/errors"
+
 	ledger "github.com/ethanfrey/ledger"
 
 	crypto "github.com/tendermint/go-crypto"
@@ -48,9 +50,13 @@ type PrivKeyLedger struct {
 	pubKey crypto.PubKey
 }
 
-func NewPrivKeyLedger() crypto.PrivKey {
+func NewPrivKeyLedger() (crypto.PrivKey, error) {
 	var pk PrivKeyLedger
-	return pk.Wrap()
+	// getPubKey will cache the pubkey for later use,
+	// this allows us to return an error early if the ledger
+	// is not plugged in
+	_, err := pk.getPubKey()
+	return pk.Wrap(), err
 }
 
 // AssertIsPrivKeyInner fulfils PrivKey Interface
@@ -74,14 +80,39 @@ func (pk *PrivKeyLedger) Sign(msg []byte) crypto.Signature {
 		panic(err)
 	}
 
-	pk.pubKey = pub
+	// if we have no pubkey yet, store it for future queries
+	if pk.pubKey.Empty() {
+		pk.pubKey = pub
+	}
 	return sig
 }
 
 // PubKey returns the stored PubKey
 // TODO: query the ledger if not there, once it is not volatile
 func (pk *PrivKeyLedger) PubKey() crypto.PubKey {
-	return pk.pubKey
+	key, err := pk.getPubKey()
+	if err != nil {
+		panic(err)
+	}
+	return key
+}
+
+// getPubKey reads the pubkey from cache or from the ledger itself
+// since this involves IO, it may return an error, which is not exposed
+// in the PubKey interface, so this function allows better error handling
+func (pk *PrivKeyLedger) getPubKey() (key crypto.PubKey, err error) {
+	// if we have no pubkey, set it
+	if pk.pubKey.Empty() {
+		dev, err := getLedger()
+		if err != nil {
+			return key, errors.WithMessage(err, "Can't connect to ledger")
+		}
+		pk.pubKey, _, err = signLedger(dev, []byte{0})
+		if err != nil {
+			return key, errors.WithMessage(err, "Can't sign with app")
+		}
+	}
+	return pk.pubKey, nil
 }
 
 // Equals fulfils PrivKey Interface
diff --git a/nano/keys_test.go b/nano/keys_test.go
index 3f2eb9a1f..3a4d9c1eb 100644
--- a/nano/keys_test.go
+++ b/nano/keys_test.go
@@ -74,16 +74,18 @@ func TestLedgerKeys(t *testing.T) {
 }
 
 func TestRealLedger(t *testing.T) {
+	assert, require := assert.New(t), require.New(t)
+
 	if os.Getenv("WITH_LEDGER") == "" {
 		t.Skip("Set WITH_LEDGER to run code on real ledger")
 	}
-
-	priv := NewPrivKeyLedger()
 	msg := []byte("kuhehfeohg")
 
-	sig := priv.Sign(msg)
+	priv, err := NewPrivKeyLedger()
+	require.Nil(err, "%+v", err)
 	pub := priv.PubKey()
+	sig := priv.Sign(msg)
 
 	valid := pub.VerifyBytes(msg, sig)
-	assert.True(t, valid)
+	assert.True(valid)
 }