Check null pointer when use --replace-restrict option

tobefuturer · Sep 24, 2016 · c9ac040 · c9ac040
1 parent 62b56fe
commit c9ac040
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 4 deletions.
diff --git 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
@@ -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.