Add API to bulk add segments, use it in elf view

5 files changed, 67 insertions, 12 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 8a217952..884861bb 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -6489,6 +6489,8 @@ namespace BinaryNinja {
 
 		/*! Add an analysis segment that specifies how data from the raw file is mapped into a virtual address space
 
+			Note that the segment added may have different size attributes than requested
+
 			\param start Starting virtual address
 			\param length Length within the virtual address space
 			\param dataOffset Data offset in the raw file
@@ -6497,6 +6499,14 @@ namespace BinaryNinja {
 		*/
 		void AddAutoSegment(uint64_t start, uint64_t length, uint64_t dataOffset, uint64_t dataLength, uint32_t flags);
 
+		/*! Add analysis segments that specify how data from the raw file is mapped into a virtual address space
+
+			\param segments Segments to add to the BinaryView
+
+			Note that the segments added may have different size attributes than requested
+		*/
+		void AddAutoSegments(const std::vector<BNSegmentInfo>& segments);
+
 		/*! Removes an automatically generated segment from the current segment mapping
 
 			\warning This action is not persistent across saving of a BNDB and must be re-applied each time a BNDB is loaded.
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 93f97dd1..f2083006 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -3445,6 +3445,14 @@ extern "C"
 		BuiltinWcscpy
 	} BNBuiltinType;
 
+	typedef struct BNSegmentInfo {
+		uint64_t start;
+		uint64_t length;
+		uint64_t dataOffset;
+		uint64_t dataLength;
+		uint32_t flags;
+	} BNSegmentInfo;
+
 	typedef bool(*BNProgressFunction)(void*, size_t, size_t);
 	typedef bool(*BNCollaborationAnalysisConflictHandler)(void*, const char** keys, BNAnalysisMergeConflict** conflicts, size_t conflictCount);
 	typedef bool(*BNCollaborationNameChangesetFunction)(void*, BNCollaborationChangeset*);
@@ -4058,6 +4066,7 @@ extern "C"
 	BINARYNINJACOREAPI bool BNPerformSearch(const char* query, const uint8_t* buffer, size_t size, bool(*callback)(void*, size_t, size_t), void* context);
 
 	BINARYNINJACOREAPI void BNAddAutoSegment(BNBinaryView* view, uint64_t start, uint64_t length, uint64_t dataOffset, uint64_t dataLength, uint32_t flags);
+	BINARYNINJACOREAPI void BNAddAutoSegments(BNBinaryView* view, const BNSegmentInfo* segmentInfo, size_t count);
 	BINARYNINJACOREAPI void BNRemoveAutoSegment(BNBinaryView* view, uint64_t start, uint64_t length);
 	BINARYNINJACOREAPI void BNAddUserSegment(BNBinaryView* view, uint64_t start, uint64_t length, uint64_t dataOffset, uint64_t dataLength, uint32_t flags);
 	BINARYNINJACOREAPI void BNRemoveUserSegment(BNBinaryView* view, uint64_t start, uint64_t length);
diff --git a/binaryview.cpp b/binaryview.cpp
index d2dbee56..7bd9cec8 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -4733,6 +4733,12 @@ void BinaryView::AddAutoSegment(
 }
 
 
+void BinaryView::AddAutoSegments(const vector<BNSegmentInfo>& segments)
+{
+	BNAddAutoSegments(m_object, segments.data(), segments.size());
+}
+
+
 void BinaryView::RemoveAutoSegment(uint64_t start, uint64_t length)
 {
 	BNRemoveAutoSegment(m_object, start, length);
diff --git a/python/binaryview.py b/python/binaryview.py
index 60cc4f28..3d5187d6 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -9055,8 +9055,22 @@ to a the type "tagRECT" found in the typelibrary "winX64common"
 		return value.value
 
 	def add_auto_segment(self, start: int, length: int, data_offset: int, data_length: int, flags: SegmentFlag) -> None:
+		"""
+		``add_auto_segment`` Adds an analysis segment that specifies how data from the raw file is mapped into a virtual address space
+
+		Note that the segments added may have different size attributes than requested
+		"""
 		core.BNAddAutoSegment(self.handle, start, length, data_offset, data_length, flags)
 
+	def add_auto_segments(self, segments: List[core.BNSegmentInfo]) -> None:
+		"""
+		``add_auto_segments`` Adds analysis segments that specify how data from the raw file is mapped into a virtual address space
+
+		Note that the segments added may have different size attributes than requested
+		"""
+		segment_arr = (core.BNSegmentInfo * len(segments))(*segments)
+		core.BNAddAutoSegments(self.handle, segment_arr, len(segments))
+
 	def remove_auto_segment(self, start: int, length: int) -> None:
 		"""
 		``remove_auto_segment`` removes an automatically generated segment from the current segment mapping.
diff --git a/view/elf/elfview.cpp b/view/elf/elfview.cpp
index 4edb2ca7..8be478f1 100644
--- a/view/elf/elfview.cpp
+++ b/view/elf/elfview.cpp
@@ -496,21 +496,37 @@ bool ElfView::Init()
 
 	m_entryPoint = m_entryPoint + imageBaseAdjustment;
 
-	for (auto& i : m_programHeaders)
+	// Add segments in bulk first
 	{
-		uint64_t adjustedVirtualAddr = i.virtualAddress + imageBaseAdjustment;
-
-		if (i.type == ELF_PT_LOAD) // || i.type == ELF_PT_GNU_RELRO)
+		vector<BNSegmentInfo> segmentsToAdd;
+		for (auto& i : m_programHeaders)
 		{
-			uint32_t flags = 0;
-			if (i.flags & 1)
-				flags |= SegmentExecutable;
-			if (i.flags & 2)
-				flags |= SegmentWritable;
-			if (i.flags & 4)
-				flags |= SegmentReadable;
-			AddAutoSegment(adjustedVirtualAddr, i.memorySize, i.offset, i.fileSize, flags);
+			uint64_t adjustedVirtualAddr = i.virtualAddress + imageBaseAdjustment;
+			if (i.type == ELF_PT_LOAD) // || i.type == ELF_PT_GNU_RELRO)
+			{
+				uint32_t flags = 0;
+				if (i.flags & 1)
+					flags |= SegmentExecutable;
+				if (i.flags & 2)
+					flags |= SegmentWritable;
+				if (i.flags & 4)
+					flags |= SegmentReadable;
+
+				BNSegmentInfo segmentInfo;
+				segmentInfo.start = adjustedVirtualAddr;
+				segmentInfo.length = i.memorySize;
+				segmentInfo.dataOffset = i.offset;
+				segmentInfo.dataLength = i.fileSize;
+				segmentInfo.flags = flags;
+				segmentsToAdd.push_back(segmentInfo);
+			}
 		}
+		AddAutoSegments(segmentsToAdd);
+	}
+
+	for (auto& i : m_programHeaders)
+	{
+		uint64_t adjustedVirtualAddr = i.virtualAddress + imageBaseAdjustment;
 
 		// Create sections for the program headers with the standard section names. This will ensure that
 		// the standard name for sections such as .dynamic will always refer to what the loader actually uses.