From a436da6c194174fa0f3b14a241ae5e7e7cddb80e Mon Sep 17 00:00:00 2001
From: Ivan Kuznetsov <me@jeiwan.ru>
Date: Sun, 10 Sep 2017 13:53:14 +0700
Subject: [PATCH] Implement ValidateAddress

---
 cli_createblockchain.go |  8 +++++++-
 cli_getbalance.go       |  8 +++++++-
 cli_send.go             | 12 +++++++++++-
 wallet.go               | 37 +++++++++++++++++++++++++------------
 4 files changed, 50 insertions(+), 15 deletions(-)

diff --git a/cli_createblockchain.go b/cli_createblockchain.go
index c2acb48..06ff1b1 100644
--- a/cli_createblockchain.go
+++ b/cli_createblockchain.go
@@ -1,8 +1,14 @@
 package main
 
-import "fmt"
+import (
+	"fmt"
+	"log"
+)
 
 func (cli *CLI) createBlockchain(address string) {
+	if !ValidateAddress(address) {
+		log.Panic("ERROR: Address is not valid")
+	}
 	bc := CreateBlockchain(address)
 	bc.db.Close()
 	fmt.Println("Done!")
diff --git a/cli_getbalance.go b/cli_getbalance.go
index 76ab0d0..27cb55e 100644
--- a/cli_getbalance.go
+++ b/cli_getbalance.go
@@ -1,8 +1,14 @@
 package main
 
-import "fmt"
+import (
+	"fmt"
+	"log"
+)
 
 func (cli *CLI) getBalance(address string) {
+	if !ValidateAddress(address) {
+		log.Panic("ERROR: Address is not valid")
+	}
 	bc := NewBlockchain(address)
 	defer bc.db.Close()
 
diff --git a/cli_send.go b/cli_send.go
index ac3e001..fb8117c 100644
--- a/cli_send.go
+++ b/cli_send.go
@@ -1,8 +1,18 @@
 package main
 
-import "fmt"
+import (
+	"fmt"
+	"log"
+)
 
 func (cli *CLI) send(from, to string, amount int) {
+	if !ValidateAddress(from) {
+		log.Panic("ERROR: Sender address is not valid")
+	}
+	if !ValidateAddress(to) {
+		log.Panic("ERROR: Recipient address is not valid")
+	}
+
 	bc := NewBlockchain(from)
 	defer bc.db.Close()
 
diff --git a/wallet.go b/wallet.go
index d527988..31a2258 100644
--- a/wallet.go
+++ b/wallet.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bytes"
 	"crypto/ecdsa"
 	"crypto/elliptic"
 	"crypto/rand"
@@ -12,6 +13,7 @@ import (
 
 const version = byte(0x00)
 const walletFile = "wallet.dat"
+const addressChecksumLen = 4
 
 // Wallet stores private and public keys
 type Wallet struct {
@@ -40,17 +42,6 @@ func (w Wallet) GetAddress() []byte {
 	return address
 }
 
-func newKeyPair() (ecdsa.PrivateKey, []byte) {
-	curve := elliptic.P256()
-	private, err := ecdsa.GenerateKey(curve, rand.Reader)
-	if err != nil {
-		log.Panic(err)
-	}
-	pubKey := append(private.PublicKey.X.Bytes(), private.PublicKey.Y.Bytes()...)
-
-	return *private, pubKey
-}
-
 // HashPubKey hashes public key
 func HashPubKey(pubKey []byte) []byte {
 	publicSHA256 := sha256.Sum256(pubKey)
@@ -65,10 +56,32 @@ func HashPubKey(pubKey []byte) []byte {
 	return publicRIPEMD160
 }
 
+// ValidateAddress check if address if valid
+func ValidateAddress(address string) bool {
+	pubKeyHash := Base58Decode([]byte(address))
+	actualChecksum := pubKeyHash[len(pubKeyHash)-addressChecksumLen:]
+	version := pubKeyHash[0]
+	pubKeyHash = pubKeyHash[1 : len(pubKeyHash)-addressChecksumLen]
+	targetChecksum := checksum(append([]byte{version}, pubKeyHash...))
+
+	return bytes.Compare(actualChecksum, targetChecksum) == 0
+}
+
 // Checksum generates a checksum for a public key
 func checksum(payload []byte) []byte {
 	firstSHA := sha256.Sum256(payload)
 	secondSHA := sha256.Sum256(firstSHA[:])
 
-	return secondSHA[:4]
+	return secondSHA[:addressChecksumLen]
+}
+
+func newKeyPair() (ecdsa.PrivateKey, []byte) {
+	curve := elliptic.P256()
+	private, err := ecdsa.GenerateKey(curve, rand.Reader)
+	if err != nil {
+		log.Panic(err)
+	}
+	pubKey := append(private.PublicKey.X.Bytes(), private.PublicKey.Y.Bytes()...)
+
+	return *private, pubKey
 }