From c2d26d423525461743b187fb7b71a04465b4a0c3 Mon Sep 17 00:00:00 2001
From: cryptoriver <114266151+cryptoriver@users.noreply.github.com>
Date: Wed, 25 Oct 2023 08:05:14 -0700
Subject: [PATCH] [fix] support checking complex method signature (#224)

* support checking complex method signature

* Add test for edge case

---------

Co-authored-by: Paul Lange <palango@gmx.de>
---
 airgap/method_registry.go      | 50 +++++++++++++++++++++++++++-------
 airgap/method_registry_test.go |  6 ++++
 2 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/airgap/method_registry.go b/airgap/method_registry.go
index 3b9f7934..5b41b203 100644
--- a/airgap/method_registry.go
+++ b/airgap/method_registry.go
@@ -77,23 +77,53 @@ func unregisteredMethodFromString(methodSignature string) (*CeloMethod, error) {
 }
 
 func validateMethodSignature(methodSig string) error {
-	// Check if the method signature contains both opening and closing parentheses
 	openParenIndex := strings.Index(methodSig, "(")
-	closeParenIndex := strings.Index(methodSig, ")")
-	if openParenIndex == -1 || closeParenIndex == -1 || openParenIndex > closeParenIndex {
+	if openParenIndex == -1 {
 		return fmt.Errorf("Invalid method signature: %s", methodSig)
 	}
 
-	// Extract the contents inside the parentheses
-	paramString := methodSig[openParenIndex+1 : closeParenIndex]
+	// Check if the method signature has non-empty method name
+	methodName := methodSig[:openParenIndex]
+	if len(methodName) == 0 {
+		return fmt.Errorf("Invalid method signature: %s", methodSig)
+	}
 
-	// If there are no contents, the signature is valid
-	if paramString == "" {
-		return nil
+	// Perform parentheses check
+	paramString := methodSig[openParenIndex:]
+	var stack []rune
+	pairs := map[rune]rune{')': '(', ']': '['}
+	for _, char := range paramString {
+		if char == '(' || char == '[' {
+			stack = append(stack, char)
+		} else if char == ')' || char == ']' {
+			if len(stack) == 0 || stack[len(stack)-1] != pairs[char] {
+				return fmt.Errorf("Invalid method signature: %s", methodSig)
+			}
+			stack = stack[:len(stack)-1]
+		}
+	}
+	if len(stack) != 0 {
+		return fmt.Errorf("Invalid method signature: %s", methodSig)
 	}
 
-	// Split the contents by comma to get individual type strings
-	methodTypes := strings.Split(paramString, ",")
+	// Extract method parameter types into a string array
+	paramString = strings.Replace(paramString, "(", " ", -1)
+	paramString = strings.Replace(paramString, ")", " ", -1)
+	paramString = strings.Replace(paramString, "[", " ", -1)
+	paramString = strings.Replace(paramString, "]", " ", -1)
+	paramString = strings.Replace(paramString, ",", " ", -1)
+
+	var methodTypes []string
+	for _, methodType := range strings.Split(paramString, " ") {
+		if methodType != "" {
+			methodTypes = append(methodTypes, methodType)
+		}
+	}
+
+	// If there are no contents, the signature is valid (contract call without arguments).
+	if len(methodTypes) == 0 {
+		return nil
+	}
 
 	// Iterate through each type string and validate
 	for _, v := range methodTypes {
diff --git a/airgap/method_registry_test.go b/airgap/method_registry_test.go
index b54ba4aa..46d6fffa 100644
--- a/airgap/method_registry_test.go
+++ b/airgap/method_registry_test.go
@@ -22,12 +22,18 @@ func Test_validateMethodSignature(t *testing.T) {
 		methodSig string
 		wantErr   bool
 	}{
+		{name: "no method name", methodSig: "()", wantErr: true},
 		{name: "valid signature with no args", methodSig: "noArgs()", wantErr: false},
 		{name: "valid signature with one arg", methodSig: "deploy(address)", wantErr: false},
 		{name: "valid signature with multiple args", methodSig: "deploy(address,uint8,bytes16,address)", wantErr: false},
+		{name: "valid signature with nested args", methodSig: "batchTransfer((address,(address,(address,uint256)[])[])[],uint256)", wantErr: false},
 		{name: "signature with invalid arg type", methodSig: "batchTransfer(DepositWalletTransfer[])", wantErr: true},
 		{name: "closing parenthesis only", methodSig: "noArgs)", wantErr: true},
 		{name: "open parenthesis only", methodSig: "noArgs(", wantErr: true},
+		{name: "missing closing bracket in the args", methodSig: "batchTransfer(bytes[)", wantErr: true},
+		{name: "mismatch parenthesis in the args", methodSig: "batchTransfer(bytes[))", wantErr: true},
+		{name: "missing open bracket in the args", methodSig: "batchTransfer(bytes])", wantErr: true},
+		{name: "missing closing bracket in the nested args", methodSig: "batchTransfer((address,(address,(address,uint256)[)[])[],uint256)", wantErr: true},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {