From 35814a605f65aafc95e8945b65e170e7ef53490b Mon Sep 17 00:00:00 2001 From: Alejo Acosta Date: Tue, 1 Oct 2024 15:55:34 -0300 Subject: [PATCH] implement pubkey comp/decompress in TxIn proto enc/decoding --- core/types/utxo.go | 50 +++++++++++++++- core/types/utxo_test.go | 126 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 174 insertions(+), 2 deletions(-) diff --git a/core/types/utxo.go b/core/types/utxo.go index 4daf7d6f49..0ff0cd272c 100644 --- a/core/types/utxo.go +++ b/core/types/utxo.go @@ -7,6 +7,7 @@ import ( "math/big" "github.com/dominant-strategies/go-quai/common" + "github.com/dominant-strategies/go-quai/crypto" "github.com/dominant-strategies/go-quai/params" ) @@ -84,7 +85,13 @@ func (txIn TxIn) ProtoEncode() (*ProtoTxIn, error) { return nil, err } - protoTxIn.PubKey = txIn.PubKey + // Compress the public key if it's uncompressed + compressedPubKey, err := compressPubKeyIfNeeded(txIn.PubKey) + if err != nil { + return nil, err + } + protoTxIn.PubKey = compressedPubKey + return protoTxIn, nil } @@ -93,10 +100,49 @@ func (txIn *TxIn) ProtoDecode(protoTxIn *ProtoTxIn) error { if err != nil { return err } - txIn.PubKey = protoTxIn.PubKey + + // Decompress the public key if it's compressed + pubKey, err := decompressPubKeyIfNeeded(protoTxIn.PubKey) + if err != nil { + return err + } + txIn.PubKey = pubKey + return nil } +// decompressPubKeyIfNeeded decompresses the public key if it's in compressed format +func decompressPubKeyIfNeeded(pubKey []byte) ([]byte, error) { + switch len(pubKey) { + case 33: // Compressed public key + uncompressedPubKey, err := crypto.DecompressPubkey(pubKey) + if err != nil { + return nil, err + } + return crypto.FromECDSAPub(uncompressedPubKey), nil + case 65: // Uncompressed public key + return pubKey, nil + default: + return nil, errors.New("invalid public key length") + } +} + +// compressPubKeyIfNeeded compresses the public key if it's in uncompressed format +func compressPubKeyIfNeeded(pubKey []byte) ([]byte, error) { + switch len(pubKey) { + case 65: // Uncompressed public key + uncompressedPubKey, err := crypto.UnmarshalPubkey(pubKey) + if err != nil { + return nil, err + } + return crypto.CompressPubkey(uncompressedPubKey), nil + case 33: // Already compressed public key + return pubKey, nil + default: + return nil, errors.New("invalid public key length") + } +} + // OutPoint defines a Qi data type that is used to track previous outputs type OutPoint struct { TxHash common.Hash `json:"txHash"` diff --git a/core/types/utxo_test.go b/core/types/utxo_test.go index c8f25cc495..c7aa4727be 100644 --- a/core/types/utxo_test.go +++ b/core/types/utxo_test.go @@ -744,3 +744,129 @@ func TestModifySSig(t *testing.T) { modifiedSigHex := hex.EncodeToString(signatureBytes) t.Log("Modified Signature: ", modifiedSigHex) } + +func TestTxInProtoDecode(t *testing.T) { + tests := []struct { + name string + pubKeyHex string + wantPubKey string + shouldFail bool + }{ + { + name: "Uncompressed public key", + pubKeyHex: "0450495cb2f9535c684ebe4687b501c0d41a623d68c118b8dcecd393370f1d90e65c4c6c44cd3fe809b41dfac9060ad84cb57e2d575fad24d25a7efa3396e73c10", + wantPubKey: "0450495cb2f9535c684ebe4687b501c0d41a623d68c118b8dcecd393370f1d90e65c4c6c44cd3fe809b41dfac9060ad84cb57e2d575fad24d25a7efa3396e73c10", + shouldFail: false, + }, + { + name: "Compressed public key", + pubKeyHex: "0250495cb2f9535c684ebe4687b501c0d41a623d68c118b8dcecd393370f1d90e6", + wantPubKey: "0450495cb2f9535c684ebe4687b501c0d41a623d68c118b8dcecd393370f1d90e65c4c6c44cd3fe809b41dfac9060ad84cb57e2d575fad24d25a7efa3396e73c10", + shouldFail: false, + }, + { + name: "Invalid public key (32 bytes)", + pubKeyHex: "50495cb2f9535c684ebe4687b501c0d41a623d68c118b8dcecd393370f1d90e6", + wantPubKey: "", + shouldFail: true, + }, + { + name: "Invalid public key (64 bytes)", + pubKeyHex: "50495cb2f9535c684ebe4687b501c0d41a623d68c118b8dcecd393370f1d90e650495cb2f9535c684ebe4687b501c0d41a623d68c118b8dcecd393370f1d90e6", + wantPubKey: "", + shouldFail: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + pubKey, _ := hex.DecodeString(tt.pubKeyHex) + protoTxIn := &ProtoTxIn{ + PreviousOutPoint: &ProtoOutPoint{ + Hash: &common.ProtoHash{Value: make([]byte, 32)}, + Index: new(uint32), + }, + PubKey: pubKey, + } + + txIn := &TxIn{} + err := txIn.ProtoDecode(protoTxIn) + + if tt.shouldFail { + if err == nil { + t.Errorf("Expected an error for invalid public key, but got none") + } + } else { + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + + gotPubKey := hex.EncodeToString(txIn.PubKey) + if gotPubKey != tt.wantPubKey { + t.Errorf("Unexpected public key. Got: %s, Want: %s", gotPubKey, tt.wantPubKey) + } + } + }) + } +} + +func TestTxInProtoEncode(t *testing.T) { + tests := []struct { + name string + pubKeyHex string + wantPubKey string + shouldFail bool + }{ + { + name: "Uncompressed public key", + pubKeyHex: "0450495cb2f9535c684ebe4687b501c0d41a623d68c118b8dcecd393370f1d90e65c4c6c44cd3fe809b41dfac9060ad84cb57e2d575fad24d25a7efa3396e73c10", + wantPubKey: "0250495cb2f9535c684ebe4687b501c0d41a623d68c118b8dcecd393370f1d90e6", + shouldFail: false, + }, + { + name: "Compressed public key", + pubKeyHex: "0250495cb2f9535c684ebe4687b501c0d41a623d68c118b8dcecd393370f1d90e6", + wantPubKey: "0250495cb2f9535c684ebe4687b501c0d41a623d68c118b8dcecd393370f1d90e6", + shouldFail: false, + }, + { + name: "Invalid public key (32 bytes)", + pubKeyHex: "50495cb2f9535c684ebe4687b501c0d41a623d68c118b8dcecd393370f1d90e6", + wantPubKey: "", + shouldFail: true, + }, + { + name: "Invalid public key (64 bytes)", + pubKeyHex: "50495cb2f9535c684ebe4687b501c0d41a623d68c118b8dcecd393370f1d90e650495cb2f9535c684ebe4687b501c0d41a623d68c118b8dcecd393370f1d90e6", + wantPubKey: "", + shouldFail: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + pubKey, _ := hex.DecodeString(tt.pubKeyHex) + txIn := &TxIn{ + PubKey: pubKey, + } + + protoTxIn, err := txIn.ProtoEncode() + if tt.shouldFail { + if err == nil { + t.Errorf("Expected an error for invalid public key, but got none") + } + } else { + if err != nil { + t.Fatalf("Unexpected error: %v", err) + } + gotPubKey := hex.EncodeToString(protoTxIn.PubKey) + if gotPubKey != tt.wantPubKey { + t.Errorf("Unexpected public key. Got: %s, Want: %s", gotPubKey, tt.wantPubKey) + } + + } + + }) + } + +}