[fix] support checking complex method signature (#224)

* support checking complex method signature * Add test for edge case --------- Co-authored-by: Paul Lange <[email protected]>
celo-org · Oct 25, 2023 · c2d26d4 · c2d26d4
1 parent 82172f9
commit c2d26d4
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 10 deletions.
diff --git 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
@@ -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) {