diff --git a/SPLPing/ICMPHeader.h b/SPLPing/ICMPHeader.h index f67673c..12832eb 100644 --- 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 index 414499b..648996c 100644 --- 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; }