From cc5e6773442a3bd71037dc2b16dfc15a020d9d5b Mon Sep 17 00:00:00 2001
From: Aurora Gaffney <aurora@blinklabs.io>
Date: Tue, 8 Jul 2025 16:29:30 -0400
Subject: [PATCH] fix: CBOR encoding for TX inputs

Fixes #1057

Signed-off-by: Aurora Gaffney <aurora@blinklabs.io>
---
 ledger/conway/conway.go   | 27 +++++++++++++++++++++++++--
 ledger/shelley/shelley.go |  4 ++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/ledger/conway/conway.go b/ledger/conway/conway.go
index f32d3dbf..c9d1f1b2 100644
--- a/ledger/conway/conway.go
+++ b/ledger/conway/conway.go
@@ -15,6 +15,7 @@
 package conway
 
 import (
+	"errors"
 	"fmt"
 
 	"github.com/blinklabs-io/gouroboros/cbor"
@@ -275,7 +276,8 @@ func (w ConwayTransactionWitnessSet) Redeemers() common.TransactionWitnessRedeem
 }
 
 type ConwayTransactionInputSet struct {
-	items []shelley.ShelleyTransactionInput
+	useSet bool
+	items  []shelley.ShelleyTransactionInput
 }
 
 func NewConwayTransactionInputSet(
@@ -288,7 +290,16 @@ func NewConwayTransactionInputSet(
 }
 
 func (s *ConwayTransactionInputSet) UnmarshalCBOR(data []byte) error {
-	// This overrides the Shelley behavior that explicitly disallowed tag-wrapped sets
+	// Check if the set is wrapped in a CBOR tag
+	// This is mostly needed so we can remember whether it was Set-wrapped for CBOR encoding
+	var tmpTag cbor.RawTag
+	if _, err := cbor.Decode(data, &tmpTag); err == nil {
+		if tmpTag.Number != cbor.CborTagSet {
+			return errors.New("unexpected tag type")
+		}
+		data = []byte(tmpTag.Content)
+		s.useSet = true
+	}
 	var tmpData []shelley.ShelleyTransactionInput
 	if _, err := cbor.Decode(data, &tmpData); err != nil {
 		return err
@@ -297,6 +308,18 @@ func (s *ConwayTransactionInputSet) UnmarshalCBOR(data []byte) error {
 	return nil
 }
 
+func (s *ConwayTransactionInputSet) MarshalCBOR() ([]byte, error) {
+	tmpItems := make([]any, len(s.items))
+	for i, item := range s.items {
+		tmpItems[i] = item
+	}
+	var tmpData any = tmpItems
+	if s.useSet {
+		tmpData = cbor.Set(tmpItems)
+	}
+	return cbor.Encode(tmpData)
+}
+
 func (s *ConwayTransactionInputSet) Items() []shelley.ShelleyTransactionInput {
 	return s.items
 }
diff --git a/ledger/shelley/shelley.go b/ledger/shelley/shelley.go
index 4cc4ac5d..5439bb46 100644
--- a/ledger/shelley/shelley.go
+++ b/ledger/shelley/shelley.go
@@ -316,6 +316,10 @@ func (s *ShelleyTransactionInputSet) UnmarshalCBOR(data []byte) error {
 	return nil
 }
 
+func (s *ShelleyTransactionInputSet) MarshalCBOR() ([]byte, error) {
+	return cbor.Encode(s.items)
+}
+
 func (s *ShelleyTransactionInputSet) Items() []ShelleyTransactionInput {
 	return s.items
 }