From d91a6c902bb2cc825449938a996d079dd201f8fa Mon Sep 17 00:00:00 2001
From: Valentin D'Emmanuele <valentin@devling.xyz>
Date: Mon, 25 Mar 2024 18:23:07 +0100
Subject: [PATCH] Fix Marshaling of Dialogue

---
 codec_test.go   | 38 ++++++++++++++++++++++++++++++++++++++
 dialogue-pdu.go |  3 ++-
 dialogue.go     | 34 +++++++++++++++++++++-------------
 tcap.go         |  1 +
 4 files changed, 62 insertions(+), 14 deletions(-)

diff --git a/codec_test.go b/codec_test.go
index fe02d1f..45572af 100644
--- a/codec_test.go
+++ b/codec_test.go
@@ -288,6 +288,44 @@ var testcases = []struct {
 			return v, nil
 		},
 	}, {
+		description: "Dialogue/AARQ/UserInformation",
+		structured: tcap.NewDialogue(
+			1, 1, // OID, Version
+			tcap.NewAARQ(
+				// Version, Context, ContextVersion
+				1, tcap.AnyTimeInfoEnquiryContext, 3,
+				tcap.NewIE(0xBE, []byte{0xde, 0xad, 0xbe, 0xef}),
+			),
+			[]byte{0xde, 0xad, 0xbe, 0xef},
+		),
+		serialized: []byte{
+			0x6b, 0x28, 0x28, 0x26, 0x06, 0x07, 0x00, 0x11, 0x86, 0x05, 0x01, 0x01, 0x01, 0xa0, 0x17, 0x60,
+			0x15, 0x80, 0x02, 0x07, 0x80, 0xa1, 0x09, 0x06, 0x07, 0x04, 0x00, 0x00, 0x01, 0x00, 0x1d, 0x03,
+			0xbe, 0x04, 0xde, 0xad, 0xbe, 0xef, 0xde, 0xad, 0xbe, 0xef,
+		},
+		parseFunc: func(b []byte) (serializable, error) {
+			v, err := tcap.ParseDialogue(b)
+			if err != nil {
+				return nil, err
+			}
+
+			b, err = v.MarshalBinary()
+			if err != nil {
+				return nil, err
+			}
+
+			// Purposely do not clean v.SingleAsn1Type.Value the first time
+
+			v, err = tcap.ParseDialogue(b)
+			if err != nil {
+				return nil, err
+			}
+
+			v.SingleAsn1Type.Value = nil
+
+			return v, nil
+		},
+	}, {
 		description: "Dialogue/AARE",
 		structured: tcap.NewDialogue(
 			1, 1, // OID, Version
diff --git a/dialogue-pdu.go b/dialogue-pdu.go
index 4c097f8..18299bf 100644
--- a/dialogue-pdu.go
+++ b/dialogue-pdu.go
@@ -700,7 +700,7 @@ func (d *DialoguePDU) ContextVersion() string {
 
 // String returns DialoguePDU in human readable string.
 func (d *DialoguePDU) String() string {
-	return fmt.Sprintf("{Type: %#x, Length: %d, ProtocolVersion: %v, ApplicationContextName: %v, Result: %v, ResultSourceDiagnostic: %v, AbortSource: %v}",
+	return fmt.Sprintf("{Type: %#x, Length: %d, ProtocolVersion: %v, ApplicationContextName: %v, Result: %v, ResultSourceDiagnostic: %v, AbortSource: %v, UserInformation: %v}",
 		d.Type,
 		d.Length,
 		d.ProtocolVersion,
@@ -708,5 +708,6 @@ func (d *DialoguePDU) String() string {
 		d.Result,
 		d.ResultSourceDiagnostic,
 		d.AbortSource,
+		d.UserInformation,
 	)
 }
diff --git a/dialogue.go b/dialogue.go
index 0af0967..bada6f7 100644
--- a/dialogue.go
+++ b/dialogue.go
@@ -76,21 +76,26 @@ func (d *Dialogue) MarshalTo(b []byte) error {
 		offset += field.MarshalLen()
 	}
 
-	if field := d.SingleAsn1Type; field != nil {
-		if err := field.MarshalTo(b[offset : offset+field.MarshalLen()]); err != nil {
-			return err
-		}
-		offset += field.MarshalLen()
+	if d.SingleAsn1Type == nil {
+		copy(b[offset:], d.Payload)
+		return nil
 	}
 
 	if field := d.DialoguePDU; field != nil {
-		if err := field.MarshalTo(b[offset : offset+field.MarshalLen()]); err != nil {
+		d.SingleAsn1Type.Value = make([]byte, field.MarshalLen())
+		if err := field.MarshalTo(d.SingleAsn1Type.Value); err != nil {
 			return err
 		}
-		offset += field.MarshalLen()
 	}
 
+	d.SingleAsn1Type.SetLength()
+	if err := d.SingleAsn1Type.MarshalTo(b[offset : offset+d.SingleAsn1Type.MarshalLen()]); err != nil {
+		return err
+	}
+	offset += d.SingleAsn1Type.MarshalLen()
+
 	copy(b[offset:], d.Payload)
+
 	return nil
 }
 
@@ -188,19 +193,22 @@ func (d *Dialogue) MarshalLen() int {
 	if field := d.ObjectIdentifier; field != nil {
 		l += field.MarshalLen()
 	}
-	if field := d.SingleAsn1Type; field != nil {
-		l += field.MarshalLen()
-	}
 	if field := d.DialoguePDU; field != nil {
-		l += field.MarshalLen()
+		l += field.MarshalLen() + 2 // 2 = singleAsn1Type IE Header
 	}
-	l += len(d.Payload)
 
-	return l
+	return l + len(d.Payload)
 }
 
 // SetLength sets the length in Length field.
 func (d *Dialogue) SetLength() {
+	if d.ObjectIdentifier != nil {
+		d.ObjectIdentifier.SetLength()
+	}
+	if d.DialoguePDU != nil {
+		d.DialoguePDU.SetLength()
+	}
+
 	d.Length = uint8(d.MarshalLen() - 2)
 	d.ExternalLength = uint8(d.MarshalLen() - 4)
 }
diff --git a/tcap.go b/tcap.go
index 3bd6ceb..b261f37 100644
--- a/tcap.go
+++ b/tcap.go
@@ -114,6 +114,7 @@ func Parse(b []byte) (*TCAP, error) {
 	if err := t.UnmarshalBinary(b); err != nil {
 		return nil, err
 	}
+
 	return t, nil
 }