Basic IsProxy() check for contract listener (#28)

* Introducing basic check for if contract is proxy
* More unit tests

contracts/contract_listener.go

@@ -16,6 +16,9 @@ type ContractListener struct {
 	*parser.BaseSolidityParserListener                        // BaseSolidityParserListener is the base listener from the Solidity parser.
 	parser                             *parser.SolidityParser // parser is the Solidity parser instance.
 	contractInfo                       ContractInfo           // contractInfo is the contract information extracted from the listener.
+	hasProxyModifier                   bool
+	hasAddressStateVariable            bool
+	hasDelegateCall                    bool
 }
 
 // NewContractListener creates a new ContractListener. It takes a SolidityParser as an argument.
@@ -77,6 +80,74 @@ func (l *ContractListener) EnterUsingDirective(ctx *parser.UsingDirectiveContext
 	}
 }
 
+// EnterFunctionDefinition is called when the parser enters a function definition.
+// It checks if there are any modifiers in the function definition and if there's a proxy modifier.
+func (l *ContractListener) EnterFunctionDefinition(ctx *parser.FunctionDefinitionContext) {
+	if ctx.AllModifierInvocation() != nil {
+		for _, modifier := range ctx.AllModifierInvocation() {
+			if strings.Contains(modifier.GetText(), "proxy") {
+				l.hasProxyModifier = true
+			}
+		}
+	}
+}
+
+// EnterFallbackFunctionDefinition is called when the parser enters a fallback function definition.
+// It checks if there's a delegatecall in the fallback function. We use this later on to determine
+// if the contract is a proxy.
+func (l *ContractListener) EnterFallbackFunctionDefinition(ctx *parser.FallbackFunctionDefinitionContext) {
+	if strings.Contains(ctx.GetText(), "delegatecall") {
+		l.hasDelegateCall = true
+	}
+}
+
+// EnterReceiveFunctionDefinition is called when the parser enters a receive function definition.
+// It checks if there's a delegatecall in the receive function. We use this later on to determine
+// if the contract is a proxy.
+func (l *ContractListener) EnterReceiveFunctionDefinition(ctx *parser.ReceiveFunctionDefinitionContext) {
+	if strings.Contains(ctx.GetText(), "delegatecall") {
+		l.hasDelegateCall = true
+	}
+}
+
+// EnterStateVariableDeclaration is called when the parser enters a state variable declaration.
+// It checks if there's a state variable that could be storing the implementation address.
+// We use this later on to determine if the contract is a proxy.
+func (l *ContractListener) EnterStateVariableDeclaration(ctx *parser.StateVariableDeclarationContext) {
+	// Check if there's a state variable that could be storing the implementation address
+	if strings.Contains(ctx.GetText(), "address") {
+		l.hasAddressStateVariable = true
+	}
+}
+
+// ExitContractDefinition is called when the parser exits a contract definition.
+// It checks if the contract is a proxy and sets the IsProxy field to true if it is.
+// It also sets the ProxyConfidence field based on current dummy algorithm.
+func (l *ContractListener) ExitContractDefinition(ctx *parser.ContractDefinitionContext) {
+	if l.hasProxyModifier && l.hasAddressStateVariable && l.hasDelegateCall {
+		l.contractInfo.IsProxy = true
+		l.contractInfo.ProxyConfidence = 100
+	} else if l.hasProxyModifier && l.hasAddressStateVariable {
+		l.contractInfo.IsProxy = true
+		l.contractInfo.ProxyConfidence = 50
+	} else if l.hasDelegateCall {
+		l.contractInfo.IsProxy = true
+		l.contractInfo.ProxyConfidence = 50
+	} else if l.hasProxyModifier {
+		l.contractInfo.IsProxy = true
+		l.contractInfo.ProxyConfidence = 25
+	}
+
+	// Following exception is if there are imports from openzeppelin that we are sure about
+	// that are proxies, but the listener doesn't detect them as proxies.
+	for _, imp := range l.contractInfo.Imports {
+		if strings.Contains(imp, "openzeppelin/contracts-upgradeable") {
+			l.contractInfo.IsProxy = true
+			l.contractInfo.ProxyConfidence = 100
+		}
+	}
+}
+
 // GetLicense returns the SPDX license identifier, if present.
 func (l *ContractListener) GetLicense() string {
 	return l.contractInfo.License
@@ -107,6 +178,16 @@ func (l *ContractListener) GetComments() []string {
 	return l.contractInfo.Comments
 }
 
+// GetIsProxy returns true if the contract is a proxy, false otherwise.
+func (l *ContractListener) GetIsProxy() bool {
+	return l.contractInfo.IsProxy
+}
+
+// GetProxyConfidence returns the confidence of the proxy detection algorithm.
+func (l *ContractListener) GetProxyConfidence() int16 {
+	return l.contractInfo.ProxyConfidence
+}
+
 // GetInfoForTests returns a map of all information extracted from the contract.
 // This is used for testing purposes only
 func (l *ContractListener) ToStruct() ContractInfo {

contracts/contract_listener_test.go

@@ -32,6 +32,75 @@ func TestContractListener(t *testing.T) {
 			contract: tests.ReadContractFileForTest(t, "Empty").Content,
 			expected: ContractInfo{},
 		},
+		{
+			name:     "Contract Without Interfaces",
+			contract: tests.ReadContractFileForTest(t, "NoInterfaces").Content,
+			expected: ContractInfo{
+				Name: "NoInterfaces",
+				Pragmas: []string{
+					"solidity ^0.8.5",
+				},
+				License: "MIT",
+			},
+		},
+		{
+			name:     "Contract Without Imports",
+			contract: tests.ReadContractFileForTest(t, "NoImports").Content,
+			expected: ContractInfo{
+				Name: "NoImports",
+				Pragmas: []string{
+					"solidity ^0.8.5",
+				},
+				License: "MIT",
+			},
+		},
+		{
+			name:     "Contract Without Pragmas",
+			contract: tests.ReadContractFileForTest(t, "NoPragmas").Content,
+			expected: ContractInfo{
+				Name:    "NoPragmas",
+				License: "MIT",
+			},
+		},
+		{
+			name:     "Contract With Single-Line Comment",
+			contract: tests.ReadContractFileForTest(t, "SingleLineComment").Content,
+			expected: ContractInfo{
+				Comments: []string{
+					"// This is a single-line comment",
+				},
+				Name: "SingleLineComment",
+				Pragmas: []string{
+					"solidity ^0.8.5",
+				},
+				License: "MIT",
+			},
+		},
+		{
+			name:     "Contract With Multi-Line Comment",
+			contract: tests.ReadContractFileForTest(t, "MultiLineComment").Content,
+			expected: ContractInfo{
+				Comments: []string{
+					"/* This is a\n multi-line comment */",
+				},
+				Name: "MultiLineComment",
+				Pragmas: []string{
+					"solidity ^0.8.5",
+				},
+				License: "MIT",
+			},
+		},
+		{
+			name:     "Contract With Different SPDX License Identifier",
+			contract: tests.ReadContractFileForTest(t, "DifferentLicense").Content,
+			expected: ContractInfo{
+				License: "GPL-3.0",
+				Name:    "DifferentLicense",
+				Pragmas: []string{
+					"solidity ^0.8.5",
+				},
+			},
+		},
 		{
 			name:     "Dummy Contract",
 			contract: tests.ReadContractFileForTest(t, "Dummy").Content,
@@ -83,6 +152,8 @@ func TestContractListener(t *testing.T) {
 					"PausableUpgradeable",
 					"SafeERC20Upgradeable",
 				},
+				IsProxy:         true,
+				ProxyConfidence: 100,
 			},
 		},
 	}

contracts/doc.go

@@ -10,6 +10,6 @@
 // Additionally, it can detect SPDX license identifiers if present.
 //
 // The extracted contract information can be accessed using the provided getter methods, which return various aspects
-// of the contract, such as the contract name, implemented interfaces, imported contracts, pragmas, comments, and SPDX license.
+// of the contract, such as the contract name, implemented interfaces, imported contracts, pragmas, comments, is proxy, and SPDX license.
 // The ContractListener also provides a convenience method, ToStruct, that returns a struct containing all the extracted information.
 package contracts

contracts/types.go

@@ -2,10 +2,12 @@ package contracts
 
 // ContractInfo contains information about a contract
 type ContractInfo struct {
-	Comments   []string `json:"comments"`   // Comments associated with the contract
-	License    string   `json:"license"`    // License information of the contract
-	Pragmas    []string `json:"pragmas"`    // Pragmas specified in the contract
-	Imports    []string `json:"imports"`    // Imported dependencies of the contract
-	Name       string   `json:"name"`       // Name of the contract
-	Implements []string `json:"implements"` // Interfaces implemented by the contract
+	Comments        []string `json:"comments"`         // Comments associated with the contract
+	License         string   `json:"license"`          // License information of the contract
+	Pragmas         []string `json:"pragmas"`          // Pragmas specified in the contract
+	Imports         []string `json:"imports"`          // Imported dependencies of the contract
+	Name            string   `json:"name"`             // Name of the contract
+	Implements      []string `json:"implements"`       // Interfaces implemented by the contract
+	IsProxy         bool     `json:"is_proxy"`         // Whether the contract is a proxy
+	ProxyConfidence int16    `json:"proxy_confidence"` // Confidence in the proxy detection
 }

data/tests/DifferentLicense.sol

@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: GPL-3.0
+pragma solidity ^0.8.5;
+
+contract DifferentLicense {
+    uint256 public value;
+}

data/tests/MultiLineComment.sol

@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.5;
+
+/* This is a
+ multi-line comment */
+contract MultiLineComment {
+    uint256 public value;
+}

data/tests/NoImports.sol

@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.5;
+
+contract NoImports {
+    uint256 public value;
+}

data/tests/NoInterfaces.sol

@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.5;
+
+contract NoInterfaces {
+    uint256 public value;
+}

data/tests/NoPragmas.sol

@@ -0,0 +1,5 @@
+// SPDX-License-Identifier: MIT
+
+contract NoPragmas {
+    uint256 public value;
+}

data/tests/SingleLineComment.sol

@@ -0,0 +1,7 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.5;
+
+// This is a single-line comment
+contract SingleLineComment {
+    uint256 public value;
+}