From 54bdad99f156d112310192c77d804c22b3988c3f Mon Sep 17 00:00:00 2001
From: anon <anon@local>
Date: Sat, 28 Dec 2024 17:25:50 -0300
Subject: [PATCH] check checksum in NewFromBase58

---
 pkg/address/address.go | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/pkg/address/address.go b/pkg/address/address.go
index b925343..45645d6 100644
--- a/pkg/address/address.go
+++ b/pkg/address/address.go
@@ -1,6 +1,7 @@
 package address
 
 import (
+	"bytes"
 	"crypto/ecdsa"
 	"crypto/sha256"
 	"encoding/hex"
@@ -77,16 +78,22 @@ func NewFromHex(s string) (Address, error) {
 }
 
 func NewFromBase58(s string) (Address, error) {
-	bytes, err := base58.Decode(s)
+	b, err := base58.Decode(s)
 	if err != nil {
 		return Address{}, err
 	}
 
-	if len(bytes) != 25 {
-		return Address{}, fmt.Errorf("tron address is not 25 bytes but %d", len(bytes))
+	if len(b) != 25 {
+		return Address{}, fmt.Errorf("tron address is not 25 bytes but %d", len(b))
 	}
 
-	return *(*[21]byte)(bytes[:21]), nil
+	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) {