Improves ICMPExtractResponseFromData failures

OliverLetterer · May 19, 2015 · 97a3a32 · 97a3a32
1 parent 36ca534
commit 97a3a32
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 9 deletions.
diff --git a/SPLPing/ICMPHeader.h b/SPLPing/ICMPHeader.h
@@ -151,6 +151,12 @@ static inline BOOL ICMPExtractResponseFromData(NSData *data, NSData **ipHeaderDa
     assert(ipHeader->protocol == 1);                               // ICMP
     size_t ipHeaderLength = (ipHeader->versionAndHeaderLength & 0x0F) * sizeof(uint32_t);
 
+    *ipHeaderData = [buffer subdataWithRange:NSMakeRange(0, sizeof(IPHeader))];
+
+    if (buffer.length >= sizeof(IPHeader) + ipHeaderLength) {
+        *ipData = [buffer subdataWithRange:NSMakeRange(sizeof(IPHeader), ipHeaderLength)];
+    }
+
     if (buffer.length < ipHeaderLength + sizeof(ICMPHeader)) {
         return NO;
     }
@@ -163,13 +169,14 @@ static inline BOOL ICMPExtractResponseFromData(NSData *data, NSData **ipHeaderDa
     uint16_t calculatedChecksum = in_cksum(icmpHeader, buffer.length - icmpHeaderOffset);
     icmpHeader->checksum = receivedChecksum;
 
-    assert(receivedChecksum == calculatedChecksum);
+    if (receivedChecksum != calculatedChecksum) {
+        NSLog(@"%s: invalid ICMP header. Checksums did not match", __PRETTY_FUNCTION__);
+        return NO;
+    }
 
-    *ipHeaderData = [buffer subdataWithRange:NSMakeRange(0, sizeof(IPHeader))];
-    *ipData = [buffer subdataWithRange:NSMakeRange(sizeof(IPHeader), ipHeaderLength)];
     *icmpHeaderData = [buffer subdataWithRange:NSMakeRange(icmpHeaderOffset, sizeof(ICMPHeader))];
     *icmpData = [buffer subdataWithRange:NSMakeRange(icmpHeaderOffset + sizeof(ICMPHeader), buffer.length - (icmpHeaderOffset + sizeof(ICMPHeader)))];
-
+    
     return YES;
 }
 

diff --git a/SPLPing/SPLPing.m b/SPLPing/SPLPing.m
@@ -208,8 +208,22 @@ - (void)socket:(CFSocketRef)socket didReadData:(NSData *)data
     NSParameterAssert([NSThread currentThread].isMainThread);
     NS_VALID_UNTIL_END_OF_SCOPE id strongSelf = self;
 
-    NSData *ipHeaderData = nil, *ipData = nil, *icmpHeaderData = nil, *icmpData = nil;
+    __block NSData *ipHeaderData = nil, *ipData = nil, *icmpHeaderData = nil, *icmpData = nil;
+    NSString *(^extractIPAddress)(void) = ^NSString *(void) {
+        if (ipHeaderData == nil) {
+            return nil;
+        }
+
+        const IPHeader *ipHeader = ipHeaderData.bytes;
+        const uint8_t *sourceAddress = ipHeader->sourceAddress;
+        return [NSString stringWithFormat:@"%d.%d.%d.%d", sourceAddress[0], sourceAddress[1], sourceAddress[2], sourceAddress[3]];
+    };
+
     if (!ICMPExtractResponseFromData(data, &ipHeaderData, &ipData, &icmpHeaderData, &icmpData)) {
+        if (ipHeaderData != nil && ![self.ip isEqualToString:extractIPAddress()]) {
+            return;
+        }
+
         NSError *error = [NSError errorWithDomain:NSURLErrorDomain code:NSURLErrorCannotDecodeRawData userInfo:@{}];
         SPLPingResponse *response = [[SPLPingResponse alloc] initWithSequenceNumber:self.currentSequenceNumber duration:[[NSDate date] timeIntervalSinceDate:self.currentStartDate] identifier:self.identifier ipAddress:nil error:error];
 
@@ -220,10 +234,7 @@ - (void)socket:(CFSocketRef)socket didReadData:(NSData *)data
         return [self _scheduleNextPing];
     }
 
-    const IPHeader *ipHeader = ipHeaderData.bytes;
-    const uint8_t *sourceAddress = ipHeader->sourceAddress;
-    NSString *ipString = [NSString stringWithFormat:@"%d.%d.%d.%d", sourceAddress[0], sourceAddress[1], sourceAddress[2], sourceAddress[3]];
-
+    NSString *ipString = extractIPAddress();
     if (![self.ip isEqualToString:ipString]) {
         return;
     }