From f90a4b95049737be6b09e3580b5f993241b10ad5 Mon Sep 17 00:00:00 2001
From: Aurora Gaffney <aurora@blinklabs.io>
Date: Tue, 11 Jun 2024 17:44:38 -0500
Subject: [PATCH] feat: custom CBOR marshaling for TX outputs (#655)

This adds custom marshaling functions to various TX output types to
match the existing custom unmarshal functions

Fixes #654
---
 ledger/alonzo.go  | 17 +++++++++++++++--
 ledger/babbage.go | 11 +++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/ledger/alonzo.go b/ledger/alonzo.go
index 889281b3..5c86d0bb 100644
--- a/ledger/alonzo.go
+++ b/ledger/alonzo.go
@@ -169,23 +169,36 @@ type AlonzoTransactionOutput struct {
 	OutputAddress     Address
 	OutputAmount      MaryTransactionOutputValue
 	TxOutputDatumHash *Blake2b256
+	legacyOutput      bool
 }
 
 func (o *AlonzoTransactionOutput) UnmarshalCBOR(cborData []byte) error {
 	// Save original CBOR
 	o.SetCbor(cborData)
-	// Try to parse as Mary output first
+	// Try to parse as legacy Mary output first
 	var tmpOutput MaryTransactionOutput
 	if _, err := cbor.Decode(cborData, &tmpOutput); err == nil {
-		// Copy from temp Shelley output to Alonzo format
+		// Copy from temp Mary output to Alonzo format
 		o.OutputAddress = tmpOutput.OutputAddress
 		o.OutputAmount = tmpOutput.OutputAmount
+		o.legacyOutput = true
 	} else {
 		return cbor.DecodeGeneric(cborData, o)
 	}
 	return nil
 }
 
+func (o *AlonzoTransactionOutput) MarshalCBOR() ([]byte, error) {
+	if o.legacyOutput {
+		tmpOutput := MaryTransactionOutput{
+			OutputAddress: o.OutputAddress,
+			OutputAmount:  o.OutputAmount,
+		}
+		return cbor.Encode(&tmpOutput)
+	}
+	return cbor.EncodeGeneric(o)
+}
+
 func (o AlonzoTransactionOutput) MarshalJSON() ([]byte, error) {
 	tmpObj := struct {
 		Address   Address                           `json:"address"`
diff --git a/ledger/babbage.go b/ledger/babbage.go
index 57f822f8..ee659aa7 100644
--- a/ledger/babbage.go
+++ b/ledger/babbage.go
@@ -342,6 +342,17 @@ func (o *BabbageTransactionOutput) UnmarshalCBOR(cborData []byte) error {
 	return nil
 }
 
+func (o *BabbageTransactionOutput) MarshalCBOR() ([]byte, error) {
+	if o.legacyOutput {
+		tmpOutput := AlonzoTransactionOutput{
+			OutputAddress: o.OutputAddress,
+			OutputAmount:  o.OutputAmount,
+		}
+		return cbor.Encode(&tmpOutput)
+	}
+	return cbor.EncodeGeneric(o)
+}
+
 func (o BabbageTransactionOutput) MarshalJSON() ([]byte, error) {
 	tmpObj := struct {
 		Address   Address                           `json:"address"`