Merge pull request #303 from blinklabs-io/feat/ledger-address-stake-address

feat: return stake address from address

Author: agaffney

ledger/common.go

@@ -202,34 +202,68 @@ type Address struct {
 	networkId      uint8
 	paymentAddress []byte
 	stakingAddress []byte
-	hrp            string
 }
 
-func (a *Address) UnmarshalCBOR(data []byte) error {
-	// Decode bytes from CBOR
-	tmpData := []byte{}
-	if _, err := cbor.Decode(data, &tmpData); err != nil {
-		return err
+// NewAddress returns an Address based on the provided bech32 address string
+func NewAddress(addr string) (Address, error) {
+	_, data, err := bech32.DecodeNoLimit(addr)
+	if err != nil {
+		return Address{}, err
 	}
+	decoded, err := bech32.ConvertBits(data, 5, 8, false)
+	if err != nil {
+		return Address{}, err
+	}
+	a := Address{}
+	a.populateFromBytes(decoded)
+	return a, nil
+}
+
+func (a *Address) populateFromBytes(data []byte) {
 	// Extract header info
-	header := tmpData[0]
+	header := data[0]
 	a.addressType = (header & addressHeaderTypeMask) >> 4
 	a.networkId = header & addressHeaderNetworkMask
 	// Extract payload
 	// NOTE: this is probably incorrect for Byron
-	payload := tmpData[1:]
+	payload := data[1:]
 	a.paymentAddress = payload[:addressHashSize]
 	a.stakingAddress = payload[addressHashSize:]
-	// Generate human readable part of address for output
-	a.hrp = a.generateHRP()
+	// Adjust stake addresses
+	if a.addressType == addressTypeNoneKey || a.addressType == addressTypeNoneScript {
+		a.stakingAddress = a.paymentAddress[:]
+		a.paymentAddress = make([]byte, 0)
+	}
+}
+
+func (a *Address) UnmarshalCBOR(data []byte) error {
+	// Decode bytes from CBOR
+	tmpData := []byte{}
+	if _, err := cbor.Decode(data, &tmpData); err != nil {
+		return err
+	}
+	a.populateFromBytes(tmpData)
 	return nil
 }
 
 func (a *Address) MarshalCBOR() ([]byte, error) {
 	return cbor.Encode(a.Bytes())
 }
 
-func (a *Address) generateHRP() string {
+// StakeAddress returns a new Address with only the stake key portion. This will return nil if the address is not a payment/staking key pair
+func (a *Address) StakeAddress() *Address {
+	if a.addressType != addressTypeKeyKey {
+		return nil
+	}
+	newAddr := &Address{
+		addressType:    addressTypeNoneKey,
+		networkId:      a.networkId,
+		stakingAddress: a.stakingAddress[:],
+	}
+	return newAddr
+}
+
+func (a Address) generateHRP() string {
 	var ret string
 	if a.addressType == addressTypeNoneKey || a.addressType == addressTypeNoneScript {
 		ret = "stake"
@@ -243,14 +277,16 @@ func (a *Address) generateHRP() string {
 	return ret
 }
 
-func (a *Address) Bytes() []byte {
+// Bytes returns the underlying bytes for the address
+func (a Address) Bytes() []byte {
 	ret := []byte{}
 	ret = append(ret, (byte(a.addressType)<<4)|(byte(a.networkId)&addressHeaderNetworkMask))
 	ret = append(ret, a.paymentAddress...)
 	ret = append(ret, a.stakingAddress...)
 	return ret
 }
 
+// String returns the bech32-encoded version of the address
 func (a Address) String() string {
 	data := a.Bytes()
 	if a.addressType == addressTypeByron {
@@ -263,7 +299,9 @@ func (a Address) String() string {
 		if err != nil {
 			panic(fmt.Sprintf("unexpected error converting data to base32: %s", err))
 		}
-		encoded, err := bech32.Encode(a.hrp, convData)
+		// Generate human readable part of address for output
+		hrp := a.generateHRP()
+		encoded, err := bech32.Encode(hrp, convData)
 		if err != nil {
 			panic(fmt.Sprintf("unexpected error encoding data as bech32: %s", err))
 		}

ledger/common_test.go

@@ -103,3 +103,63 @@ func TestMultiAssetJson(t *testing.T) {
 		}
 	}
 }
+
+func TestAddressFromBytes(t *testing.T) {
+	testDefs := []struct {
+		addressBytesHex string
+		expectedAddress string
+	}{
+		{
+			addressBytesHex: "11e1317b152faac13426e6a83e06ff88a4d62cce3c1634ab0a5ec1330952563c5410bff6a0d43ccebb7c37e1f69f5eb260552521adff33b9c2",
+			expectedAddress: "addr1z8snz7c4974vzdpxu65ruphl3zjdvtxw8strf2c2tmqnxz2j2c79gy9l76sdg0xwhd7r0c0kna0tycz4y5s6mlenh8pq0xmsha",
+		},
+		{
+			addressBytesHex: "013f35615835258addded1c2e169f3a2ab4ae94d606bde030e7947f5184ff5f8e3d43ce6b19ec4197e331e86d0f5e58b02d7a75b5e74cff95d",
+			expectedAddress: "addr1qyln2c2cx5jc4hw768pwz60n5245462dvp4auqcw09rl2xz07huw84puu6cea3qe0ce3apks7hjckqkh5ad4uax0l9ws0q9xty",
+		},
+		{
+			addressBytesHex: "7121bd8c2e0df2fbe92137f78dbaba48f62308e52303049f0d628b6c4c",
+			expectedAddress: "addr1wysmmrpwphe0h6fpxlmcmw46frmzxz89yvpsf8cdv29kcnqsw3vw6",
+		},
+		{
+			addressBytesHex: "61cfe224295a282d69edda5fa8de4f131e2b9cd21a6c9235597fa4ff6b",
+			expectedAddress: "addr1v887yfpftg5z660dmf063hj0zv0zh8xjrfkfyd2e07j076cecha5k",
+		},
+	}
+	for _, testDef := range testDefs {
+		addr := Address{}
+		addr.populateFromBytes(test.DecodeHexString(testDef.addressBytesHex))
+		if addr.String() != testDef.expectedAddress {
+			t.Fatalf("address did not match expected value, got: %s, wanted: %s", addr.String(), testDef.expectedAddress)
+		}
+	}
+}
+
+func TestAddressStakeAddress(t *testing.T) {
+	testDefs := []struct {
+		address              string
+		expectedStakeAddress string
+	}{
+		{
+			address:              "addr1q8fv95d4g2599v3gzq7wnva34ykt4d2zerl0wyke36zml0neqj84x95mgp694rv8gfqy6u67ms38lx30texma843yd5qmvkqcz",
+			expectedStakeAddress: "stake1u9usfr6nz6d5qaz63kr5yszdwd0dcgnlngh4und7n6cjx6qh02h9m",
+		},
+		{
+			address:              "addr1q8uas4shlrnxhd8dnqxesk9hlgmtx65xlmkq9c7acfa2ksvjn4k4w8d7lwnzkf3dkq26kxz3re50h89adduskx08rr6qq2g0f2",
+			expectedStakeAddress: "stake1uxff6m2hrkl0hf3tyckmq9dtrpg3u68mnj7kk7gtr8n33aq5fqskz",
+		},
+		{
+			address:              "addr1q9h4f2vhh5vnqgnsejan3psw6mj3a504fxlqm2eh3262qufesdvfs83ulr22vprsv9mwnt0vgkfwxlflxkns32twqzdqjpq2na",
+			expectedStakeAddress: "stake1uyucxkycrc70349xq3cxzahf4hkytyhr05lntfcg49hqpxsqlrayk",
+		},
+	}
+	for _, testDef := range testDefs {
+		addr, err := NewAddress(testDef.address)
+		if err != nil {
+			t.Fatalf("failed to decode address: %s", err)
+		}
+		if addr.StakeAddress().String() != testDef.expectedStakeAddress {
+			t.Fatalf("stake address did not match expected value, got: %s, wanted: %s", addr.StakeAddress().String(), testDef.expectedStakeAddress)
+		}
+	}
+}