From 0767091824792354005502377403661643401745 Mon Sep 17 00:00:00 2001
From: Glenn Smith <glenn@vector35.com>
Date: Fri, 25 Oct 2024 17:16:50 -0400
Subject: Fix not being able to set UIDF on function parameters

Mirrors the UI code's method for handling this situation: if the variable's index field is 0, then it is a function parameter and it needs to have its value set at the function's entry point.
---
 python/function.py | 50 +++++++++++++++++++++++++++++---------------------
 1 file changed, 29 insertions(+), 21 deletions(-)

(limited to 'python')

diff --git a/python/function.py b/python/function.py
index e43c11e2..8b6239f3 100644
--- a/python/function.py
+++ b/python/function.py
@@ -2828,16 +2828,20 @@ class Function:
 			>>> var_value = PossibleValueSet.constant(5)
 			>>> current_function.set_user_var_value(mlil_var, def_address, var_value)
 		"""
-		var_defs = self.mlil.get_var_definitions(var)
-		if var_defs is None:
-			raise ValueError("Could not get definition for Variable")
-		found = False
-		for site in var_defs:
-			if site.address == def_addr:
-				found = True
-				break
-		if not found:
-			raise ValueError("No definition for Variable found at given address")
+		if var.index == 0:
+			# Special case: function parameters have index 0 and are defined at the start of the function
+			def_addr = self.start
+		else:
+			var_defs = self.mlil.get_var_definitions(var)
+			if var_defs is None:
+				raise ValueError("Could not get definition for Variable")
+			found = False
+			for site in var_defs:
+				if site.address == def_addr:
+					found = True
+					break
+			if not found:
+				raise ValueError("No definition for Variable found at given address")
 		def_site = core.BNArchitectureAndAddress()
 		def_site.arch = self.arch.handle
 		def_site.address = def_addr
@@ -2852,17 +2856,21 @@ class Function:
 		:param int def_addr: Address of the definition site of the variable
 		:rtype: None
 		"""
-		var_defs = self.mlil.get_var_definitions(var)
-		if var_defs is None:
-			raise ValueError("Could not get definition for Variable")
-
-		found = False
-		for site in var_defs:
-			if site.address == def_addr:
-				found = True
-				break
-		if not found:
-			raise ValueError("No definition for Variable found at given address")
+		if var.index == 0:
+			# Special case: function parameters have index 0 and are defined at the start of the function
+			def_addr = self.start
+		else:
+			var_defs = self.mlil.get_var_definitions(var)
+			if var_defs is None:
+				raise ValueError("Could not get definition for Variable")
+
+			found = False
+			for site in var_defs:
+				if site.address == def_addr:
+					found = True
+					break
+			if not found:
+				raise ValueError("No definition for Variable found at given address")
 		def_site = core.BNArchitectureAndAddress()
 		def_site.arch = self.arch.handle
 		def_site.address = def_addr
-- 
cgit v1.3.1