gotron/pkg/address/address.go

package address

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/mr-tron/base58"
	"strings"
)

type Address [21]byte

func (a *Address) UnmarshalJSON(bytes []byte) error {
	s := ""

	if err := json.Unmarshal(bytes, &s); err != nil {
		return fmt.Errorf("unmarshal string: %w", err)
	}

	addr, err := NewFromBase58(s)
	if err != nil {
		return fmt.Errorf("parse address: %w", err)
	}

	*a = addr
	return nil
}

func (a Address) String() string {
	return a.Base58()
}

func (a Address) Bytes() []byte {
	return a[:]
}

func (a Address) Hex() string {
	return hex.EncodeToString(a[:])
}

func (a Address) Base58() string {
	var tmp []byte
	tmp = append(tmp, a[:]...)

	hash := SHA256(SHA256(tmp))
	checksum := hash[:4]
	tmp = append(tmp, checksum...)

	return base58.Encode(tmp)
}

func (a Address) ETH() common.Address {
	return common.BytesToAddress(a[1:])
}

func NewFromETH(eth common.Address) Address {
	return Address(append([]byte{65}, eth.Bytes()...))
}

func NewFromBytes(b []byte) (Address, error) {
	if len(b) != 21 {
		return Address{}, fmt.Errorf("tron address is not 21 bytes but %d", len(b))
	}
	return *(*[21]byte)(b), nil
}

func NewFromHex(s string) (Address, error) {
	addrBytes, err := hex.DecodeString(s)
	if err != nil {
		return Address{}, err
	}
	return *(*[21]byte)(addrBytes), nil
}

func NewFromBase58(s string) (Address, error) {
	b, err := base58.Decode(s)
	if err != nil {
		return Address{}, err
	}

	if len(b) != 25 {
		return Address{}, fmt.Errorf("tron address is not 25 bytes but %d", len(b))
	}

	hash := SHA256(SHA256(b[:21]))
	checksum := hash[:4]
	if !bytes.Equal(checksum, b[21:]) {
		return Address{}, fmt.Errorf("invalid checksum")
	}

	return *(*[21]byte)(b[:21]), nil
}

func NewFromString(s string) (Address, error) {
	if strings.HasPrefix(s, "T") {
		return NewFromBase58(s)
	}
	if strings.HasPrefix(s, "41") {
		return NewFromHex(s)
	}
	return [21]byte{}, fmt.Errorf("%s is not a Tron address", s)
}

func NewFromKey(k *ecdsa.PrivateKey) Address {
	return NewFromETH(crypto.PubkeyToAddress(k.PublicKey))
}

func NewFromKeyBytes(k []byte) (Address, error) {
	key, err := crypto.ToECDSA(k)
	if err != nil {
		return Address{}, err
	}
	return NewFromKey(key), nil
}

func NewFromKeyHex(k string) (Address, error) {
	keyBytes, err := hex.DecodeString(k)
	if err != nil {
		return Address{}, err
	}
	return NewFromKeyBytes(keyBytes)
}

func SHA256(s []byte) []byte {
	h := sha256.New()
	h.Write(s)
	bs := h.Sum(nil)
	return bs
}