From c9ac04015a2b624e8aabcd67e3806b4097f520d2 Mon Sep 17 00:00:00 2001
From: eugeneyang <EugeneYang@eugene-MB1.lan>
Date: Sun, 25 Sep 2016 04:17:04 +0800
Subject: [PATCH] Check null pointer when use --replace-restrict option

---
 source/RSScanMethodVisitor.m | 13 ++++++++++++-
 source/restore-symbol.m      | 16 +++++++++++++---
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/source/RSScanMethodVisitor.m b/source/RSScanMethodVisitor.m
index 75b5a18..9e6d8a4 100644
--- a/source/RSScanMethodVisitor.m
+++ b/source/RSScanMethodVisitor.m
@@ -12,6 +12,7 @@
 #import "CDLCDylib.h"
 #import "CDOCClass.h"
 #import "CDOCCategory.h"
+#import "CDOCClassReference.h"
 #import "CDOCMethod.h"
 //#import "CDTypeController.h"
 
@@ -64,8 +65,18 @@ - (void)visitClassMethod:(CDOCMethod *)method;
     if (method.address == 0 ) {
         return;
     }
-    NSString *name = [NSString stringWithFormat:@"+[%@ %@]", _context.name, method.name];
     
+    NSString *name = nil;
+    if ([_context isKindOfClass:[CDOCClass class]]) {
+        name = [NSString stringWithFormat:@"+[%@ %@]", _context.name, method.name];
+    } else if([_context isKindOfClass:[CDOCCategory class]]) {
+        NSString * className = [[(CDOCCategory *)_context classRef] className];
+        if (!className) className = @"";
+        name = [NSString stringWithFormat:@"+[%@(%@) %@]", className ,_context.name, method.name];
+    }
+    
+    if (!name) return;
+
     RSSymbol *s = [RSSymbol symbolWithName:name address:method.address];
     
     [self.collector addSymbol:s];
diff --git a/source/restore-symbol.m b/source/restore-symbol.m
index 1b817c6..2a2c291 100644
--- a/source/restore-symbol.m
+++ b/source/restore-symbol.m
@@ -151,11 +151,21 @@ void restore_symbol(NSString * inpath, NSString *outpath, NSString *jsonPath, bo
         CDLCSegment * restrict_seg = [machOFile segmentWithName:@"__RESTRICT"];
         
         struct segment_command *restrict_seg_cmd = (struct segment_command *)((char *)outData.mutableBytes + restrict_seg.commandOffset);
-        struct section *restrict_section = (struct section *)((char *)outData.mutableBytes + restrict_seg.commandOffset + (Is32Bit? sizeof(struct segment_command) : sizeof(struct segment_command_64)));
+        struct section *restrict_section = NULL;
         
+        int cmd_size = (Is32Bit? sizeof(struct segment_command) : sizeof(struct segment_command_64));
+        if (restrict_seg.cmdsize > cmd_size) {
+            restrict_section = (struct section *)((char *)outData.mutableBytes + restrict_seg.commandOffset + cmd_size);
+        }
+        
+        if (restrict_seg && restrict_section) {
+            fprintf(stderr, "rename segment __RESTRICT  in mach-o header.\n");
+            strncpy(restrict_seg_cmd -> segname, "__restrict", 16);
+            strncpy(restrict_section -> segname, "__restrict", 16);
+        } else {
+            fprintf(stderr, "No section (__RESTRICT,__restrict) in mach-o header.\n");
+        }
         
-        strncpy(restrict_seg_cmd -> segname, "__restrict", 16);
-        strncpy(restrict_section -> segname, "__restrict", 16);
     }
     
     //LC_CODE_SIGNATURE need align 16 byte, so add padding at end of string table.