From d1cfecf516cd5d5385160e8945f4017533f10952 Mon Sep 17 00:00:00 2001
From: Zichuan Li <34680029+river-li@users.noreply.github.com>
Date: Mon, 29 Jul 2024 14:36:23 -0400
Subject: [PATCH] fix: fix bug relate to defining descriptor types

When we tried to add backward compatibility on bndb files, we added a
function `get_type`, but the result of
```python
bv.get_code_refs_for_type("EFI_PEI_NOTIFY_DESCRIPTOR")
```
and
```
bv.get_code_refs_for_type(bv.types("EFI_PEI_NOTIFY_DESCRIPTOR")
```
is different.

Fixed this bug by not defining descriptor types on existing bndb
---
 ppi.py       | 18 ++++++++++--------
 protocols.py |  6 +++---
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/ppi.py b/ppi.py
index 9416586..a6e9440 100644
--- a/ppi.py
+++ b/ppi.py
@@ -138,7 +138,7 @@ def define_pei_idt(bv: BinaryView, task: BackgroundTask) -> bool:
     return True
 
 
-def _define_descriptor(bv: BinaryView, task: BackgroundTask, descriptor_type, param) -> Optional[bool]:
+def _define_descriptor(bv: BinaryView, task: BackgroundTask, descriptor_type: str, param) -> Optional[bool]:
     """
     Define the descriptor type at `param`'s address. The `param` should be a ConstPtr parameter.
     If the analysis got cancelled or encountered an error, return False. If it doesn't meet conditions, return None.
@@ -150,7 +150,7 @@ def _define_descriptor(bv: BinaryView, task: BackgroundTask, descriptor_type, pa
 
     var = bv.get_data_var_at(var_addr)
     if var:
-        if descriptor_type == var.type:
+        if descriptor_type in str(var.type):
             # already defined
             return
     sym = bv.get_symbol_at(var_addr)
@@ -172,8 +172,8 @@ def _define_descriptor(bv: BinaryView, task: BackgroundTask, descriptor_type, pa
         return
 
     # define types for guid and notify entrypoint
-    protocol_name = lookup_and_define_guid(bv, notify_descriptor["Guid"])
-    if protocol_name is False:
+    guid_name = lookup_and_define_guid(bv, notify_descriptor["Guid"])
+    if guid_name is False:
         return False
 
     if "Notify" in notify_descriptor:
@@ -181,10 +181,10 @@ def _define_descriptor(bv: BinaryView, task: BackgroundTask, descriptor_type, pa
         func = bv.get_function_at(notify_entrypoint)
         if not func:
             return
-        if not protocol_name:
+        if not guid_name:
             func_name = non_conflicting_symbol_name(bv, "UnknownNotify")
         else:
-            func_name = non_conflicting_symbol_name(bv, f"Notify{get_var_name_from_type(protocol_name)}")
+            func_name = non_conflicting_symbol_name(bv, f"Notify{get_var_name_from_type(guid_name)}")
         func.type = f"EFI_STATUS {func_name}(EFI_PEI_SERVICES **PeiServices, EFI_PEI_NOTIFY_DESCRIPTOR* NotifyDescriptor, VOID* Ppi)"
         bv.update_analysis_and_wait()
         return propagate_function_param_types(bv, task, func)
@@ -199,8 +199,10 @@ def define_pei_descriptor(bv: BinaryView, task: BackgroundTask) -> bool:
     supported yet.
     """
     descriptor_types_names = ["EFI_PEI_NOTIFY_DESCRIPTOR", "EFI_PEI_PPI_DESCRIPTOR"]
-    descriptor_types = [get_type(bv, des) for des in descriptor_types_names]
-    for descriptor_type in descriptor_types:
+    for descriptor_type in descriptor_types_names:
+        if bv.types.get(descriptor_type) is None:
+            log_warn("Cannot define descriptors in existing BNDB, continue without define descriptors")
+            return True
         refs = list(bv.get_code_refs_for_type(descriptor_type))
         for ref in refs:
             if task.cancelled:
diff --git a/protocols.py b/protocols.py
index eb4848a..86c7924 100644
--- a/protocols.py
+++ b/protocols.py
@@ -113,17 +113,17 @@ def lookup_protocol_guid(guid: bytes) -> Tuple[Optional[str], Optional[str]]:
 
 def lookup_and_define_guid(bv: BinaryView, addr: int) -> bool | Optional[str]:
     """
-    Input an address, define the guid there, lookup the protocol mapping and return the protocol name.
+    Input an address, define the guid there, lookup the protocol mapping and return the guid name
     """
     guid = bv.read(addr, 16)
     if not guid or len(guid) != 16:
         return False
-    protocol_name, guid_name = lookup_protocol_guid(guid)
+    _, guid_name = lookup_protocol_guid(guid)
     if guid_name is None:
         guid_name = non_conflicting_symbol_name(bv, "UnknownGuid")
     bv.define_user_data_var(addr, 'EFI_GUID', guid_name)
 
-    return protocol_name
+    return guid_name
 
 
 def nonconflicting_variable_name(func: Function, base_name: str) -> str: