From 44babe3baffd4d47bb43fcb09c0f948a8bc476b0 Mon Sep 17 00:00:00 2001
From: Lephenixnoir <sebastien.michelland@protonmail.com>
Date: Sun, 20 Aug 2023 18:02:27 +0200
Subject: [PATCH] lib: allow declaring exclusive claims multiple times

---
 include/fxos/disassembly.h | 11 ++++++++--
 lib/disassembly.cpp        | 44 ++++++++++++++++++++++++++++----------
 2 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/include/fxos/disassembly.h b/include/fxos/disassembly.h
index 93c45ff..617e336 100644
--- a/include/fxos/disassembly.h
+++ b/include/fxos/disassembly.h
@@ -144,6 +144,9 @@ struct Claim
     /* Utility to check for intersections */
     bool intersects(Claim const &other) const;
 
+    /* Check equality of claims (raw equality) */
+    bool operator==(Claim const &other) const;
+
     /* String representation */
     std::string str() const;
 };
@@ -194,11 +197,15 @@ struct Disassembly
     /* Access the claim that owns the address, if there is one */
     Claim const *getClaimAt(uint32_t address);
 
-    /* Access the first claim that overlaps this region, if any */
+    /* Find the first claim that overlaps this region, if any */
     Claim const *findClaimConflict(uint32_t address, int size);
 
+    /* Find all (or up to max ≥ 0) claims that overlaps this region */
+    std::vector<Claim const *> findClaimConflicts(
+        uint32_t address, int size, int max = -1);
+
     /* Add a new exclusive claim. If there is any intersection with previous
-       claims, this fails. */
+       claims which do not compare equal to c, this fails. */
     bool addExclusiveClaim(Claim const &c);
 
     // TODO: Add non-exclusive claims/handle collisions
diff --git a/lib/disassembly.cpp b/lib/disassembly.cpp
index 976b60a..6a987bd 100644
--- a/lib/disassembly.cpp
+++ b/lib/disassembly.cpp
@@ -65,6 +65,12 @@ bool Claim::intersects(Claim const &other) const
     return inter_start < inter_end;
 }
 
+bool Claim::operator==(Claim const &other) const
+{
+    return this->address == other.address && this->size == other.size
+           && this->type == other.type && this->owner == other.owner;
+}
+
 std::string Claim::str() const
 {
     std::string details = format(" (claim 0x%08x:%d)", address, size);
@@ -140,7 +146,8 @@ Function *Disassembly::getFunctionAt(uint32_t pc)
         return &it->second;
 }
 
-Claim const *Disassembly::findClaimConflict(uint32_t address, int size)
+std::vector<Claim const *> Disassembly::findClaimConflicts(
+    uint32_t address, int size, int max)
 {
     Claim fake_claim = {
         .address = address,
@@ -148,6 +155,7 @@ Claim const *Disassembly::findClaimConflict(uint32_t address, int size)
         .type = 0,
         .owner = 0,
     };
+    std::vector<Claim const *> conflicts;
 
     /* Find the first claim whose start is [> address] */
     auto it = this->claims.upper_bound(fake_claim);
@@ -155,19 +163,26 @@ Claim const *Disassembly::findClaimConflict(uint32_t address, int size)
     if(it != this->claims.begin())
         it--;
 
-    while(it != this->claims.end()) {
+    while(it != this->claims.end()
+          && (max < 0 || conflicts.size() < (size_t)max)) {
         /* We completely passed address+size, no conflict found */
         if(it->address >= address + size)
-            return nullptr;
+            break;
 
         /* There is an intersection */
         if(it->intersects(fake_claim))
-            return &*it;
+            conflicts.push_back(&*it);
 
         it++;
     }
 
-    return nullptr;
+    return conflicts;
+}
+
+Claim const *Disassembly::findClaimConflict(uint32_t address, int size)
+{
+    auto claims = findClaimConflicts(address, size, 1);
+    return claims.size() > 0 ? claims[0] : nullptr;
 }
 
 Claim const *Disassembly::getClaimAt(uint32_t address)
@@ -177,15 +192,22 @@ Claim const *Disassembly::getClaimAt(uint32_t address)
 
 bool Disassembly::addExclusiveClaim(Claim const &c)
 {
-    Claim const *conflict = this->findClaimConflict(c.address, c.size);
-    if(conflict) {
+    auto conflicts = this->findClaimConflicts(c.address, c.size);
+    bool exclusive = true;
+
+    for(auto conflict: conflicts) {
+        /* Allow declaring the same claim twice */
+        if(*conflict == c)
+            continue;
+
         FxOS_log(ERR, "exclusive claim for %s conflicts with %s", c.str(),
-            conflict->str());
-        return false;
+            conflicts[0]->str());
+        exclusive = false;
     }
 
-    this->claims.insert(c);
-    return true;
+    if(exclusive)
+        this->claims.insert(c);
+    return exclusive;
 }