From 12c6ccbe1cef0958c5b720fde1b00a5aa36c9ec0 Mon Sep 17 00:00:00 2001
From: dongin <needshield@gmail.com>
Date: Sat, 13 Oct 2018 22:30:37 +0900
Subject: [PATCH] modify for supporting changed ToByte function's return
 value(s)

---
 mocks/mock_key_service.go | 34 +++++++++++++++++++++-------------
 server/server_test.go     | 10 +++++++---
 util.go                   |  6 +++++-
 3 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/mocks/mock_key_service.go b/mocks/mock_key_service.go
index 5f2bd418..f272903c 100644
--- a/mocks/mock_key_service.go
+++ b/mocks/mock_key_service.go
@@ -164,7 +164,11 @@ func MockStoreKey(key bifrost.Key, keyDirPath string) error {
 	keyFilePath := filepath.Join(keyDirPath, key.ID())
 
 	// byte formatted key
-	keyBytes := key.ToByte()
+	keyBytes, err := key.ToByte()
+	if err != nil {
+		return err
+	}
+
 	isPrivateKey := key.IsPrivate()
 
 	// make json marshal
@@ -224,19 +228,21 @@ func newMockPriKey(internalKey *ecdsa.PrivateKey) *MockPriKey {
 func (mockPriKey *MockPriKey) ID() bifrost.KeyID {
 	// get keyBytes from key
 	mockPub := &mockPriKey.internalPriKey.PublicKey
-	pubKeyBytes := newMockPubKey(mockPub).ToByte()
+	keyBitString := elliptic.Marshal(mockPub.Curve, mockPub.X, mockPub.Y)
 
 	// get ski from keyBytes
 	hash := sha256.New()
-	hash.Write(pubKeyBytes)
-	ski := hash.Sum(nil)
+	hash.Write(keyBitString)
+	hashValue := hash.Sum(nil)
+
+	ski := make([]byte, 20)
+	ski = hashValue[:20]
 
 	return "IT" + base58.Encode(ski)
 }
 
-func (mockPriKey *MockPriKey) ToByte() []byte {
-	keyBytes, _ := x509.MarshalECPrivateKey(mockPriKey.internalPriKey)
-	return keyBytes
+func (mockPriKey *MockPriKey) ToByte() ([]byte, error) {
+	return x509.MarshalECPrivateKey(mockPriKey.internalPriKey)
 }
 
 func (mockPriKey *MockPriKey) KeyGenOpt() string {
@@ -260,19 +266,21 @@ func newMockPubKey(internalKey *ecdsa.PublicKey) *MockPubKey {
 
 func (mockPubKey *MockPubKey) ID() bifrost.KeyID {
 	// get keyBytes from key
-	keyBytes := mockPubKey.ToByte()
+	keyBitString := elliptic.Marshal(mockPubKey.internalPubKey.Curve, mockPubKey.internalPubKey.X, mockPubKey.internalPubKey.Y)
 
 	// get ski from keyBytes
 	hash := sha256.New()
-	hash.Write(keyBytes)
-	ski := hash.Sum(nil)
+	hash.Write(keyBitString)
+	hashValue := hash.Sum(nil)
+
+	ski := make([]byte, 20)
+	ski = hashValue[:20]
 
 	return "IT" + base58.Encode(ski)
 }
 
-func (mockPubKey *MockPubKey) ToByte() []byte {
-	keyBytes, _ := x509.MarshalPKIXPublicKey(mockPubKey.internalPubKey)
-	return keyBytes
+func (mockPubKey *MockPubKey) ToByte() ([]byte, error) {
+	return x509.MarshalPKIXPublicKey(mockPubKey.internalPubKey)
 }
 
 func (mockPubKey *MockPubKey) KeyGenOpt() string {
diff --git a/server/server_test.go b/server/server_test.go
index 72cd0f5e..9fb49ac5 100644
--- a/server/server_test.go
+++ b/server/server_test.go
@@ -14,10 +14,12 @@ import (
 func TestServer_validateRequestPeerInfo_whenValidPeerInfo(t *testing.T) {
 	// given
 	keyOpt := mocks.NewMockKeyOpts()
+	keyBytes, err := keyOpt.PubKey.ToByte()
+	assert.NoError(t, err)
 
 	peerInfo := &bifrost.PeerInfo{
 		IP:          "127.0.0.1",
-		PubKeyBytes: keyOpt.PubKey.ToByte(),
+		PubKeyBytes: keyBytes,
 		IsPrivate:   keyOpt.PubKey.IsPrivate(),
 	}
 
@@ -43,17 +45,19 @@ func TestServer_BifrostStream(t *testing.T) {
 	s := mocks.NewMockServer()
 
 	keyOpt := mocks.NewMockKeyOpts()
+	keyBytes, err := keyOpt.PubKey.ToByte()
+	assert.NoError(t, err)
 
 	peerInfo := &bifrost.PeerInfo{
 		IP:          "127.0.0.1",
-		PubKeyBytes: keyOpt.PubKey.ToByte(),
+		PubKeyBytes: keyBytes,
 		IsPrivate:   keyOpt.PubKey.IsPrivate(),
 	}
 
 	mockStreamServer := mocks.NewMockStreamServer(*peerInfo)
 
 	// when
-	err := s.BifrostStream(mockStreamServer)
+	err = s.BifrostStream(mockStreamServer)
 
 	// then
 	assert.NoError(t, err)
diff --git a/util.go b/util.go
index 4daea2f1..c3778705 100644
--- a/util.go
+++ b/util.go
@@ -42,7 +42,11 @@ type KeyOpts struct {
 }
 
 func BuildResponsePeerInfo(pubKey Key, metaData map[string]string) (*pb.Envelope, error) {
-	b := pubKey.ToByte()
+	b, err := pubKey.ToByte()
+
+	if err != nil {
+		return nil, err
+	}
 
 	pi := &PeerInfo{
 		PubKeyBytes: b,