From a1366dbd005ffbc83e9d2b9fd00eeda0a2e13cad Mon Sep 17 00:00:00 2001
From: Yogesh Deshpande <yogesh.deshpande@arm.com>
Date: Mon, 18 Sep 2023 07:23:15 -0400
Subject: [PATCH] More corrections

Signed-off-by: Yogesh Deshpande <yogesh.deshpande@arm.com>
---
 corim/entity_extension.go | 50 +++++++++++++++++++++++----------------
 1 file changed, 30 insertions(+), 20 deletions(-)

diff --git a/corim/entity_extension.go b/corim/entity_extension.go
index 8afbd290..9966b046 100644
--- a/corim/entity_extension.go
+++ b/corim/entity_extension.go
@@ -13,35 +13,45 @@ type EntityExtension struct {
 	Param2 byte   `cbor:"1,keyasint,omitempty" json:"param2,omitempty"`
 }
 
-func (e *EntityExtension) FromCBOR(data []byte) error {
+type ProfileXEntity struct {
+	Entity
+	Extension EntityExtension `cbor:"8,keyasint" json:"extension"`
+}
 
-	if err := dm.Unmarshal(data, e); err != nil {
-		return fmt.Errorf("Failed to unMarshal extension")
-	}
-	return nil
+func NewProfileXEntity() *ProfileXEntity {
+	return &ProfileXEntity{}
 }
 
-func (e *EntityExtension) ToCBOR() ([]byte, error) {
-	data, err := em.Marshal(e)
-	if err != nil {
-		return nil, fmt.Errorf("unable to CBOR Marshal extension")
+func (e *ProfileXEntity) SetEntityExtension(p1 string, p2 byte) {
+
+	e.Extension.Param1 = p1
+	e.Extension.Param2 = p2
+}
+
+func (o *ProfileXEntity) FromCBOR(data []byte) error {
+	if err := o.Valid(); err != nil {
+		return fmt.Errorf("invalid Profile XEntity %w", err)
 	}
-	fmt.Printf("ToCBOR a success")
-	return data, nil
+	return dm.Unmarshal(data, o)
 }
 
-func (e *EntityExtension) ExtensionToJSON() ([]byte, error) {
-	data, err := json.Marshal(e)
-	if err != nil {
-		return nil, fmt.Errorf("unable to JSON Marshal extension")
+func (o *ProfileXEntity) ToCBOR() ([]byte, error) {
+	if err := o.Valid(); err != nil {
+		return nil, fmt.Errorf("invalid Profile XEntity %w", err)
 	}
-	return data, nil
+	return em.Marshal(o)
 }
 
-func (e *EntityExtension) ExtensionFromJSON(data []byte) (interface{}, error) {
+func (o *ProfileXEntity) FromJSON(data []byte) error {
+	if err := o.Valid(); err != nil {
+		return err
+	}
+	return json.Unmarshal(data, o)
+}
 
-	if err := json.Unmarshal(data, e); err != nil {
-		return nil, fmt.Errorf("Failed to JSON UnMarshal extension")
+func (o *ProfileXEntity) ToJSON() ([]byte, error) {
+	if err := o.Valid(); err != nil {
+		return nil, err
 	}
-	return e, nil
+	return json.Marshal(o)
 }