From fee9bfd3afeef02e237e93335fa2076c30b9d6d6 Mon Sep 17 00:00:00 2001
From: Ivan Kuznetsov <me@jeiwan.ru>
Date: Mon, 29 Jan 2018 15:53:44 +0700
Subject: [PATCH] Fix address version processing in Base58 encoding/decoding

---
 base58.go      | 25 +++++++++----------------
 base58_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 16 deletions(-)
 create mode 100644 base58_test.go

diff --git a/base58.go b/base58.go
index 5db078e..38c1c35 100644
--- a/base58.go
+++ b/base58.go
@@ -22,38 +22,31 @@ func Base58Encode(input []byte) []byte {
 		result = append(result, b58Alphabet[mod.Int64()])
 	}
 
-	ReverseBytes(result)
-	for _, b := range input {
-		if b == 0x00 {
-			result = append([]byte{b58Alphabet[0]}, result...)
-		} else {
-			break
-		}
+	// https://en.bitcoin.it/wiki/Base58Check_encoding#Version_bytes
+	if input[0] == 0x00 {
+		result = append(result, b58Alphabet[0])
 	}
 
+	ReverseBytes(result)
+
 	return result
 }
 
 // Base58Decode decodes Base58-encoded data
 func Base58Decode(input []byte) []byte {
 	result := big.NewInt(0)
-	zeroBytes := 0
 
 	for _, b := range input {
-		if b == 0x00 {
-			zeroBytes++
-		}
-	}
-
-	payload := input[zeroBytes:]
-	for _, b := range payload {
 		charIndex := bytes.IndexByte(b58Alphabet, b)
 		result.Mul(result, big.NewInt(58))
 		result.Add(result, big.NewInt(int64(charIndex)))
 	}
 
 	decoded := result.Bytes()
-	decoded = append(bytes.Repeat([]byte{byte(0x00)}, zeroBytes), decoded...)
+
+	if input[0] == b58Alphabet[0] {
+		decoded = append([]byte{0x00}, decoded...)
+	}
 
 	return decoded
 }
diff --git a/base58_test.go b/base58_test.go
new file mode 100644
index 0000000..c65875b
--- /dev/null
+++ b/base58_test.go
@@ -0,0 +1,24 @@
+package main
+
+import (
+	"encoding/hex"
+	"log"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBase58(t *testing.T) {
+	rawHash := "00010966776006953D5567439E5E39F86A0D273BEED61967F6"
+	hash, err := hex.DecodeString(rawHash)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	encoded := Base58Encode(hash)
+	assert.Equal(t, "16UwLL9Risc3QfPqBUvKofHmBQ7wMtjvM", string(encoded))
+
+	decoded := Base58Decode([]byte("16UwLL9Risc3QfPqBUvKofHmBQ7wMtjvM"))
+	assert.Equal(t, strings.ToLower("00010966776006953D5567439E5E39F86A0D273BEED61967F6"), hex.EncodeToString(decoded))
+}