Add compression support to both EncodePrivateKey and DecodePrivateKey

addrconvs.go

@@ -109,7 +109,7 @@ func DecodeAddress(addr string) (addrHash []byte, net btcwire.BitcoinNet, err er
 
 // EncodePrivateKey takes a 32-byte private key and encodes it into the
 // Wallet Import Format (WIF).
-func EncodePrivateKey(privKey []byte, net btcwire.BitcoinNet) (string, error) {
+func EncodePrivateKey(privKey []byte, net btcwire.BitcoinNet, compressed bool) (string, error) {
 	if len(privKey) != 32 {
 		return "", ErrMalformedPrivateKey
 	}
@@ -125,27 +125,45 @@ func EncodePrivateKey(privKey []byte, net btcwire.BitcoinNet) (string, error) {
 	}
 
 	tosum := append([]byte{netID}, privKey...)
+	if compressed {
+		tosum = append(tosum, 0x01)
+	}
 	cksum := btcwire.DoubleSha256(tosum)
 
 	// Private key before base58 encoding is 1 byte for netID, 32 bytes for
-	// privKey, plus 4 bytes of checksum.
+	// privKey, plus an optional byte (0x01) if copressed, plus 4 bytes of checksum.
-	a := make([]byte, 37, 37)
+	encodeLen := 37
+	if compressed {
+		encodeLen += 1
+	}
+	a := make([]byte, encodeLen, encodeLen)
 	a[0] = netID
 	copy(a[1:], privKey)
-	copy(a[32+1:], cksum[:4])
+	if compressed {
-
+		copy(a[32+1:], []byte{0x01})
+		copy(a[32+1+1:], cksum[:4])
+	} else {
+		copy(a[32+1:], cksum[:4])
+	}
 	return Base58Encode(a), nil
 }
 
 // DecodePrivateKey takes a Wallet Import Format (WIF) string and
 // decodes into a 32-byte private key.
-func DecodePrivateKey(wif string) ([]byte, btcwire.BitcoinNet, error) {
+func DecodePrivateKey(wif string) ([]byte, btcwire.BitcoinNet, bool, error) {
 	decoded := Base58Decode(wif)
+	decodedLen := len(decoded)
+	compressed := false
 
-	// Length of decoded privkey must be 32 bytes + 1 byte for 0x80
+	// Length of decoded privkey must be 32 bytes + an optional 1 byte (0x01)
-	// + 4 bytes of checksum
+	// if compressed, plus 1 byte for netID + 4 bytes of checksum
-	if len(decoded) != 32+5 {
+	if decodedLen == 32+6 {
-		return nil, 0, ErrMalformedPrivateKey
+		compressed = true
+		if decoded[33] != 0x01 {
+			return nil, 0, compressed, ErrMalformedPrivateKey
+		}
+	} else if decodedLen != 32+5 {
+		return nil, 0, compressed, ErrMalformedPrivateKey
 	}
 
 	var net btcwire.BitcoinNet
@@ -155,20 +173,25 @@ func DecodePrivateKey(wif string) ([]byte, btcwire.BitcoinNet, error) {
 	case TestNetKey:
 		net = btcwire.TestNet3
 	default:
-		return nil, 0, ErrUnknownNet
+		return nil, 0, compressed, ErrUnknownNet
 	}
 
 	// Checksum is first four bytes of double SHA256 of the identifier byte
 	// and privKey.  Verify this matches the final 4 bytes of the decoded
 	// private key.
-	tosum := decoded[:32+1]
+	var tosum []byte
+	if compressed {
+		tosum = decoded[:32+1+1]
+	} else {
+		tosum = decoded[:32+1]
+	}
 	cksum := btcwire.DoubleSha256(tosum)[:4]
-	if !bytes.Equal(cksum, decoded[len(decoded)-4:]) {
+	if !bytes.Equal(cksum, decoded[decodedLen-4:]) {
-		return nil, 0, ErrMalformedPrivateKey
+		return nil, 0, compressed, ErrMalformedPrivateKey
 	}
 
 	privKey := make([]byte, 32, 32)
 	copy(privKey[:], decoded[1:32+1])
 
-	return privKey, net, nil
+	return privKey, net, compressed, nil
 }

addrconvs_test.go

@@ -12,16 +12,23 @@ import (
 )
 
 var encodePrivateKeyTests = []struct {
-	in  []byte
+	in         []byte
-	net btcwire.BitcoinNet
+	net        btcwire.BitcoinNet
-	out string
+	compressed bool
+	out        string
 }{
 	{[]byte{
 		0x0c, 0x28, 0xfc, 0xa3, 0x86, 0xc7, 0xa2, 0x27,
 		0x60, 0x0b, 0x2f, 0xe5, 0x0b, 0x7c, 0xae, 0x11,
 		0xec, 0x86, 0xd3, 0xbf, 0x1f, 0xbe, 0x47, 0x1b,
 		0xe8, 0x98, 0x27, 0xe1, 0x9d, 0x72, 0xaa, 0x1d,
-	}, btcwire.MainNet, "5HueCGU8rMjxEXxiPuD5BDku4MkFqeZyd4dZ1jvhTVqvbTLvyTJ"},
+	}, btcwire.MainNet, false, "5HueCGU8rMjxEXxiPuD5BDku4MkFqeZyd4dZ1jvhTVqvbTLvyTJ"},
+	{[]byte{
+		0xdd, 0xa3, 0x5a, 0x14, 0x88, 0xfb, 0x97, 0xb6,
+		0xeb, 0x3f, 0xe6, 0xe9, 0xef, 0x2a, 0x25, 0x81,
+		0x4e, 0x39, 0x6f, 0xb5, 0xdc, 0x29, 0x5f, 0xe9,
+		0x94, 0xb9, 0x67, 0x89, 0xb2, 0x1a, 0x03, 0x98,
+	}, btcwire.TestNet3, true, "cV1Y7ARUr9Yx7BR55nTdnR7ZXNJphZtCCMBTEZBJe1hXt2kB684q"},
 }
 
 var encodeTests = []struct {
@@ -106,10 +113,10 @@ func TestDecodeAddresses(t *testing.T) {
 }
 
 func TestEncodeDecodePrivateKey(t *testing.T) {
-	for _, test := range encodePrivateKeyTests {
+	for x, test := range encodePrivateKeyTests {
-		wif, err := btcutil.EncodePrivateKey(test.in, test.net)
+		wif, err := btcutil.EncodePrivateKey(test.in, test.net, test.compressed)
 		if err != nil {
-			t.Error(err)
+			t.Errorf("%x: %v", x, err)
 			continue
 		}
 		if wif != test.out {
@@ -118,12 +125,12 @@ func TestEncodeDecodePrivateKey(t *testing.T) {
 			continue
 		}
 
-		key, _, err := btcutil.DecodePrivateKey(test.out)
+		key, _, compressed, err := btcutil.DecodePrivateKey(test.out)
 		if err != nil {
 			t.Error(err)
 			continue
 		}
-		if !bytes.Equal(key, test.in) {
+		if !bytes.Equal(key, test.in) || compressed != test.compressed {
 			t.Errorf("TestEncodeDecodePrivateKey failed: want '%x', got '%x'",
 				test.out, key)
 		}