Thiagodeev/Fix SNIP-12 JSON Marshaling for TypedData (#789)

thiagodeev · web-flow · commit 4cf08ae7d74d · 2025-09-19T10:08:56.000-03:00
* refactor: update benchmarks in typedData_test.go

- Moved fileNames to a global variable for better accessibility across tests.
- Updated BenchmarkGetMessageHash to utilize the new fileNames variable.
- Removed unused BMockTypedData helper
- Added BenchmarkUnmarshalJSON to test the performance of the UnmarshalJSON function with various test data files.

* test: add TestMarshalJSON to validate JSON marshaling of TypedData

* refactor: add JSON tags to TypedData fields for improved JSON serialization

* feat: implement MarshalJSON for TypeDefinition to enable JSON serialization

* refactor: enhance Domain struct for improved JSON marshaling

- Added custom logic in UnmarshalJSON to handle both `chainId` and `chain_id` fields.
- Introduced flags to track the type of `chainId` for accurate marshaling.
- Implemented MarshalJSON to ensure `chainId` is serialized correctly based on its original format.

* refactor: enhance Domain struct for JSON marshaling with revision support

- Added logic in UnmarshalJSON to handle the `revision` field, including a flag to track its type.
- Updated MarshalJSON to ensure the `Revision` field is serialized correctly based on its original format.
- Improved comments for clarity on the purpose of marshaling logic for both `chainId` and `revision` fields.

* refactor: improve handling of chainId and chain_id in EncodeData function

- Updated logic to accommodate both `chainId` and `chain_id` fields in the domainMap.
- Added checks to copy values between the two fields as a workaround for compatibility.
- Enhanced comments for better understanding of the changes made.

* chore: update changelog with the changes

* refactor: improve error messages in UnmarshalJSON for Domain struct

- Updated error messages in UnmarshalJSON to specify 'JSON field' instead of 'struct' for clarity

* refactor: simplify error handling in UnmarshalJSON for Domain struct

- Removed redundant error check for 'chainId' retrieval in UnmarshalJSON.

* refactor: enhance error handling in MarshalJSON for Domain struct

- Improved error message for 'chainId' parsing to provide clearer context on failure.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -25,10 +25,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - The internal `tryUnwrapToRPCErr` func of the `rpc` pkg was renamed to `UnwrapToRPCErr` and moved to the new package.
   - The `Err` function now have a specific case for the `InternalError` code.
 
+### Fixed
+- The `typedData.TypedData` was not being marshaled exactly as it is in the original JSON. Now, the original JSON is preserved,
+  so the output of `TypedData.MarshalJSON()` is exactly as the original JSON.
+
 ### Dev updates
 - New `internal/tests/jsonrpc_spy.go` file containing a `Spy` type for spying JSON-RPC calls in tests. The
 old `rpc/spy_test.go` file was removed.
 - New `mocks/mock_client.go` file containing a mock of the `client.Client` type (`client.ClientI` interface).
+- New benchmarks and tests for the `typedData` pkg.
 
 ## [0.15.0](https://github.com/NethermindEth/starknet.go/releases/tag/v0.15.0) - 2025-09-03
 ### Changed
diff --git a/typedData/typedData.go b/typedData/typedData.go
@@ -19,18 +19,23 @@ import (
 var typeNameRegexp = regexp.MustCompile(`[^\(\),\s]+`)
 
 type TypedData struct {
-	Types       map[string]TypeDefinition
-	PrimaryType string
-	Domain      Domain
-	Message     map[string]any
-	Revision    *revision
+	Types       map[string]TypeDefinition `json:"types"`
+	PrimaryType string                    `json:"primaryType"`
+	Domain      Domain                    `json:"domain"`
+	Message     map[string]any            `json:"message"`
+	Revision    *revision                 `json:"-"`
 }
 
 type Domain struct {
 	Name     string `json:"name"`
 	Version  string `json:"version"`
 	ChainId  string `json:"chainId"`
 	Revision uint8  `json:"revision,omitempty"`
+
+	// Flags to deal with edge cases, used to marshal the `Domain` exactly as it is in the original JSON.
+	hasStringChainId  bool `json:"-"`
+	hasOldChainIdName bool `json:"-"`
+	HasStringRevision bool `json:"-"`
 }
 
 type TypeDefinition struct {
@@ -421,8 +426,16 @@ func EncodeData(typeDef *TypeDefinition, td *TypedData, context ...string) (enc
 			return enc, err
 		}
 
+		// The ChainId can be either `chainId` or `chain_id`.
 		// ref: https://community.starknet.io/t/signing-transactions-and-off-chain-messages/66
-		domainMap["chain_id"] = domainMap["chainId"]
+		// We find the one that contains the value and we copy the value to the other field.
+		// It's an workaround to handle both cases.
+
+		if domainMap["chainId"] != nil {
+			domainMap["chain_id"] = domainMap["chainId"]
+		} else {
+			domainMap["chainId"] = domainMap["chain_id"]
+		}
 
 		return encodeData(typeDef, td, domainMap, false, context...)
 	}
@@ -886,7 +899,7 @@ func (domain *Domain) UnmarshalJSON(data []byte) error {
 	getField := func(fieldName string) (string, error) {
 		value, ok := dec[fieldName]
 		if !ok {
-			return "", fmt.Errorf("error getting the value of '%s' from 'domain' struct", fieldName)
+			return "", fmt.Errorf("error getting the value of '%s' from 'domain' JSON field", fieldName)
 		}
 
 		return fmt.Sprintf("%v", value), nil
@@ -902,34 +915,125 @@ func (domain *Domain) UnmarshalJSON(data []byte) error {
 		return err
 	}
 
-	revision, err := getField("revision")
-	if err != nil {
-		revision = "0"
+	// Custom logic to handle the `revision` field,
+	// used to marshal the Revision exactly as it is in the original JSON.
+	rawRevision, ok := dec["revision"]
+	if !ok {
+		rawRevision = "0"
 	}
-	numRevision, err := strconv.ParseUint(revision, 10, 8)
-	if err != nil {
-		return err
+
+	var numRevision uint64
+	switch revision := rawRevision.(type) {
+	case string:
+		domain.HasStringRevision = true
+		numRevision, err = strconv.ParseUint(revision, 10, 8)
+		if err != nil {
+			return err
+		}
+	case float64:
+		domain.HasStringRevision = false
+		numRevision = uint64(revision)
 	}
 
-	chainId, err := getField("chainId")
-	if err != nil {
+	// Custom logic to handle the `chainId` field,
+	// used to marshal the ChainId exactly as it is in the original JSON.
+	rawChainId, ok := dec["chainId"]
+	if !ok {
+		err = errors.New("error getting the value of 'chainId' from 'domain' JSON field")
 		if numRevision == 1 {
 			return err
 		}
-		var err2 error
+
+		// `chain_id` was also used in the past, so we check for it if the `chainId` field is not found
 		// ref: https://community.starknet.io/t/signing-transactions-and-off-chain-messages/66
-		chainId, err2 = getField("chain_id")
-		if err2 != nil {
+		rawChainId, ok = dec["chain_id"]
+		if !ok {
+			err2 := errors.New("error getting the value of 'chain_id' from 'domain' JSON field")
+
 			return fmt.Errorf("%w: %w", err, err2)
 		}
+		domain.hasOldChainIdName = true
 	}
 
+	switch rawChainId.(type) {
+	case string:
+		domain.hasStringChainId = true
+	case float64:
+		domain.hasStringChainId = false
+	}
+	chainId := fmt.Sprintf("%v", rawChainId)
+
+	// Final step
 	*domain = Domain{
 		Name:     name,
 		Version:  version,
 		ChainId:  chainId,
 		Revision: uint8(numRevision),
+
+		hasStringChainId:  domain.hasStringChainId,
+		hasOldChainIdName: domain.hasOldChainIdName,
+		HasStringRevision: domain.HasStringRevision,
 	}
 
 	return nil
 }
+
+// MarshalJSON implements the json.Marshaler interface for Domain.
+// Some logic was added to marshal the `Domain` exactly as it is in the original JSON.
+func (domain Domain) MarshalJSON() ([]byte, error) {
+	var chainId any
+	var revision any
+	var err error
+
+	// E.g: if it's `1`, we will marshal it as `1`, not `"1"`.
+	if domain.hasStringChainId {
+		chainId = domain.ChainId
+	} else {
+		chainId, err = strconv.Atoi(domain.ChainId)
+		if err != nil {
+			return nil, fmt.Errorf("cannot parse 'chain_id' value: %w", err)
+		}
+	}
+
+	if domain.Revision == 0 {
+		revision = nil
+	} else {
+		if domain.HasStringRevision {
+			revision = strconv.FormatUint(uint64(domain.Revision), 10)
+		} else {
+			revision = domain.Revision
+		}
+	}
+
+	// The purpose here is to marshal the `Domain` exactly as it is in the original JSON.
+	// This is achieved by having two chainId fields, one for the old name and one for the new name,
+	// and using the `omitempty` tag to only include one of them, the one that is the same as the original JSON.
+	// Similar for the `Revision` field.
+	var temp struct {
+		Name       string `json:"name"`
+		Version    string `json:"version"`
+		ChainIdOld any    `json:"chain_id,omitempty"` // old chainId json name
+		ChainIdNew any    `json:"chainId,omitempty"`  // new chainId json name
+		Revision   any    `json:"revision,omitempty"`
+	}
+	temp.Name = domain.Name
+	temp.Version = domain.Version
+	temp.Revision = revision
+
+	if domain.hasOldChainIdName {
+		temp.ChainIdOld = chainId
+
+		return json.Marshal(temp)
+	}
+
+	temp.ChainIdNew = chainId
+
+	return json.Marshal(temp)
+}
+
+// MarshalJSON implements the json.Marshaler interface for TypeDefinition
+//
+//nolint:gocritic //  json.Marshaler interface requires a value receiver
+func (td TypeDefinition) MarshalJSON() ([]byte, error) {
+	return json.Marshal(td.Parameters)
+}
diff --git a/typedData/typedData_test.go b/typedData/typedData_test.go
@@ -10,19 +10,9 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-var typedDataExamples = make(map[string]TypedData)
-
-// TestMain initialises test data by loading TypedData examples from JSON files.
-// It reads multiple test files and stores them in the typedDataExamples map
-// before running the tests.
-//
-// Parameters:
-//   - m: The testing.M object that provides the test runner
-//
-// Returns:
-//   - None (calls os.Exit directly)
-func TestMain(m *testing.M) {
-	fileNames := []string{
+var (
+	typedDataExamples = make(map[string]TypedData)
+	fileNames         = []string{
 		"baseExample",
 		"example_array",
 		"example_baseTypes",
@@ -34,7 +24,18 @@ func TestMain(m *testing.M) {
 		"allInOne",
 		"example_enumNested",
 	}
+)
 
+// TestMain initialises test data by loading TypedData examples from JSON files.
+// It reads multiple test files and stores them in the typedDataExamples map
+// before running the tests.
+//
+// Parameters:
+//   - m: The testing.M object that provides the test runner
+//
+// Returns:
+//   - None (calls os.Exit directly)
+func TestMain(m *testing.M) {
 	for _, fileName := range fileNames {
 		var ttd TypedData
 		content, err := os.ReadFile(fmt.Sprintf("./testData/%s.json", fileName))
@@ -52,23 +53,52 @@ func TestMain(m *testing.M) {
 	os.Exit(m.Run())
 }
 
-// BMockTypedData is a helper function for benchmarks that loads a base example
-// TypedData from a JSON file.
-//
-// Parameters:
-//   - b: The testing.B object used for benchmarking
-//
-// Returns:
-//   - ttd: A TypedData instance loaded from the base example file
-func BMockTypedData(b *testing.B) (ttd TypedData) {
-	b.Helper()
-	content, err := os.ReadFile("./testData/baseExample.json")
-	require.NoError(b, err)
+// BenchmarkUnmarshalJSON is a benchmark function for testing the TypedData.UnmarshalJSON function.
+func BenchmarkUnmarshalJSON(b *testing.B) {
+	for _, fileName := range fileNames {
+		rawData, err := os.ReadFile(fmt.Sprintf("./testData/%s.json", fileName))
+		require.NoError(b, err)
 
-	err = json.Unmarshal(content, &ttd)
-	require.NoError(b, err)
+		b.Run(fileName, func(b *testing.B) {
+			for b.Loop() {
+				var ttd TypedData
+				err = json.Unmarshal(rawData, &ttd)
+				require.NoError(b, err)
+			}
+		})
+	}
+}
+
+// BenchmarkGetMessageHash is a benchmark function for testing the GetMessageHash function.
+func BenchmarkGetMessageHash(b *testing.B) {
+	addr := "0xdeadbeef"
+
+	for key, typedData := range typedDataExamples {
+		b.Run(key, func(b *testing.B) {
+			for b.Loop() {
+				_, err := typedData.GetMessageHash(addr)
+				if err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}
+
+// TestMarshalJSON tests the MarshalJSON function. It marshals the TypedData and compares the result
+// with the original raw data.
+func TestMarshalJSON(t *testing.T) {
+	for _, filename := range fileNames {
+		t.Run(filename, func(t *testing.T) {
+			rawData, err := os.ReadFile(fmt.Sprintf("./testData/%s.json", filename))
+			require.NoError(t, err)
 
-	return
+			marshaledData, err := json.Marshal(typedDataExamples[filename])
+			require.NoError(t, err)
+
+			require.JSONEq(t, string(rawData), string(marshaledData))
+		})
+	}
 }
 
 // TestMessageHash tests the GetMessageHash function.
@@ -157,29 +187,6 @@ func TestGetMessageHash(t *testing.T) {
 	}
 }
 
-// BenchmarkGetMessageHash is a benchmark function for testing the GetMessageHash function.
-//
-// It tests the performance of the GetMessageHash function by running it with different input sizes.
-// The input size is determined by the bit length of the address parameter, which is converted from
-// a hexadecimal string to a big integer using the HexToBN function from the utils package.
-//
-// Parameters:
-//   - b: a testing.B object that provides methods for benchmarking the function
-//
-// Returns:
-//
-//	none
-func BenchmarkGetMessageHash(b *testing.B) {
-	ttd := BMockTypedData(b)
-
-	addr := "0xCD2a3d9F938E13CD947Ec05AbC7FE734Df8DD826"
-	b.Run(fmt.Sprintf("input_size_%d", len(addr)), func(b *testing.B) {
-		result, err := ttd.GetMessageHash(addr)
-		require.NoError(b, err)
-		require.NotEmpty(b, result)
-	})
-}
-
 // TestGeneral_GetTypeHash tests the GetTypeHash function.
 //
 // It tests the GetTypeHash function by calling it with different input values
