implement pubkey comp/decompress in TxIn proto enc/decoding

dominant-strategies · Oct 10, 2024 · 35814a6 · 35814a6
1 parent fae5d6c
commit 35814a6
Show file tree

Hide file tree

Showing 2 changed files with 174 additions and 2 deletions.
diff --git 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
@@ -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)
+				}
+
+			}
+
+		})
+	}
+
+}