From bb573243a1a8d7d4e6587fa0cd9616c2a106280e Mon Sep 17 00:00:00 2001
From: k1kolodz <krystian.kolodziej@assecobs.pl>
Date: Tue, 8 Jun 2021 11:39:23 +0200
Subject: [PATCH] Fixed reading of "small" methods

---
 Source/CDLCSegment.h            |  2 ++
 Source/CDLCSegment.m            | 15 +++++++++++++++
 Source/CDMachOFile.h            |  1 +
 Source/CDMachOFile.m            | 33 ++++++++++++++++++++++++++++++++-
 Source/CDObjectiveC2Processor.m |  6 ++++++
 5 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/Source/CDLCSegment.h b/Source/CDLCSegment.h
index a7ac145d..7ca1c7eb 100644
--- a/Source/CDLCSegment.h
+++ b/Source/CDLCSegment.h
@@ -38,10 +38,12 @@ extern NSString *CDSegmentEncryptionTypeName(CDSegmentEncryptionType type);
 - (NSString *)flagDescription;
 
 - (BOOL)containsAddress:(NSUInteger)address;
+- (BOOL)containsOffset:(NSUInteger)address;
 - (CDSection *)sectionContainingAddress:(NSUInteger)address;
 - (CDSection *)sectionWithName:(NSString *)name;
 - (NSUInteger)fileOffsetForAddress:(NSUInteger)address;
 - (NSUInteger)segmentOffsetForAddress:(NSUInteger)address;
+- (NSUInteger)addressForDataOffset:(NSUInteger)offset;
 
 - (void)writeSectionData;
 
diff --git a/Source/CDLCSegment.m b/Source/CDLCSegment.m
index 4ce256ad..14b36594 100644
--- a/Source/CDLCSegment.m
+++ b/Source/CDLCSegment.m
@@ -90,6 +90,11 @@ - (NSUInteger)vmaddr;
     return _segmentCommand.vmaddr;
 }
 
+- (NSUInteger)vmsize;
+{
+    return _segmentCommand.vmsize;
+}
+
 - (NSUInteger)fileoff;
 {
     return _segmentCommand.fileoff;
@@ -169,6 +174,11 @@ - (BOOL)containsAddress:(NSUInteger)address;
     return (address >= _segmentCommand.vmaddr) && (address < _segmentCommand.vmaddr + _segmentCommand.vmsize);
 }
 
+- (BOOL)containsOffset:(NSUInteger)address
+{
+    return (address >= _segmentCommand.fileoff) && (address < _segmentCommand.fileoff + _segmentCommand.filesize);
+}
+
 - (CDSection *)sectionContainingAddress:(NSUInteger)address;
 {
     for (CDSection *section in self.sections) {
@@ -194,6 +204,11 @@ - (NSUInteger)fileOffsetForAddress:(NSUInteger)address;
     return [[self sectionContainingAddress:address] fileOffsetForAddress:address];
 }
 
+- (NSUInteger)addressForDataOffset:(NSUInteger)offset
+{
+    return self.vmaddr + (offset - self.fileoff);
+}
+
 - (NSUInteger)segmentOffsetForAddress:(NSUInteger)address;
 {
     return [self fileOffsetForAddress:address] - self.fileoff;
diff --git a/Source/CDMachOFile.h b/Source/CDMachOFile.h
index ac5c402b..e60d296f 100644
--- a/Source/CDMachOFile.h
+++ b/Source/CDMachOFile.h
@@ -57,6 +57,7 @@ typedef enum : NSUInteger {
 - (NSString *)stringAtAddress:(NSUInteger)address;
 
 - (NSUInteger)dataOffsetForAddress:(NSUInteger)address;
+- (NSUInteger)addressForDataOffset:(NSUInteger)offset;
 
 - (const void *)bytes;
 - (const void *)bytesAtOffset:(NSUInteger)offset;
diff --git a/Source/CDMachOFile.m b/Source/CDMachOFile.m
index 36983ffb..f3e8c0d2 100644
--- a/Source/CDMachOFile.m
+++ b/Source/CDMachOFile.m
@@ -319,6 +319,17 @@ - (CDLCSegment *)segmentContainingAddress:(NSUInteger)address;
     return nil;
 }
 
+- (CDLCSegment *)segmentContainingOffset:(NSUInteger)offset;
+{
+    for (id loadCommand in _loadCommands) {
+        if ([loadCommand isKindOfClass:[CDLCSegment class]] && [loadCommand containsOffset:offset]) {
+            return loadCommand;
+        }
+    }
+
+    return nil;
+}
+
 - (void)showWarning:(NSString *)warning;
 {
     NSLog(@"Warning: %@", warning);
@@ -356,7 +367,8 @@ - (NSString *)stringAtAddress:(NSUInteger)address;
     CDSection *section = [segment sectionContainingAddress:address];
     if ([[section sectionName] isEqualToString:@"__objc_selrefs"]) {
         const void * reference = [self.data bytes] + offset;
-        offset = ([self ptrSize] == 8) ? *((uint64_t *)reference) : *((uint32_t *)reference);
+        NSUInteger vmaddr = ([self ptrSize] == 8) ? *((uint64_t *)reference) : *((uint32_t *)reference);
+        offset = [self dataOffsetForAddress:vmaddr];
     }
 
     ptr = (uint8_t *)[self.data bytes] + offset;
@@ -383,6 +395,25 @@ - (NSUInteger)dataOffsetForAddress:(NSUInteger)address;
     return [segment fileOffsetForAddress:address];
 }
 
+- (NSUInteger)addressForDataOffset:(NSUInteger)offset
+{
+    if (offset == 0)
+        return 0;
+
+    CDLCSegment *segment = [self segmentContainingOffset:offset];
+    if (segment == nil) {
+        NSLog(@"Error: Cannot find segment for data offset 0x%08lx in segmentContainingOffset:", offset);
+        exit(5);
+    }
+
+    if ([segment isProtected]) {
+        NSLog(@"Error: Segment is protected.");
+        exit(5);
+    }
+
+    return [segment addressForDataOffset:offset];
+}
+
 - (const void *)bytes;
 {
     return [self.data bytes];
diff --git a/Source/CDObjectiveC2Processor.m b/Source/CDObjectiveC2Processor.m
index 6e6ce922..3c1141e0 100644
--- a/Source/CDObjectiveC2Processor.m
+++ b/Source/CDObjectiveC2Processor.m
@@ -423,6 +423,12 @@ - (NSArray *)loadMethodsAtAddress:(uint64_t)address extendedMethodTypesCursor:(C
             objc2Method.name  = [cursor readPtr:small];
             objc2Method.types = [cursor readPtr:small];
             objc2Method.imp   = [cursor readPtr:small];
+            
+            if (small) {
+                objc2Method.name  = [self.machOFile addressForDataOffset:objc2Method.name];
+                objc2Method.types = [self.machOFile addressForDataOffset:objc2Method.types];
+            }
+            
             NSString *name    = [self.machOFile stringAtAddress:objc2Method.name];
             NSString *types   = [self.machOFile stringAtAddress:objc2Method.types];