Add some extra checks to AMQParser

(not exhaustive) Return default values when attempting to read beyond the bounds of the NSData. In future, we should consider using subdataWithRange: et al to avoid(?) these checks, and to cover other possibilities, like length being unreadable. [#5]
rabbitmq · Mar 16, 2016 · a3a6b51 · a3a6b51
1 parent df5299d
commit a3a6b51
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 26 deletions.
diff --git a/RMQClient/AMQParser.m b/RMQClient/AMQParser.m
@@ -11,15 +11,17 @@
 @interface AMQParser ()
 @property (nonatomic, readwrite) const char *cursor;
 @property (nonatomic, readwrite) const char *end;
+@property (nonatomic, readwrite) NSData *data;
 @end
 
 @implementation AMQParser
 
 - (instancetype)initWithData:(NSData *)data {
     self = [super init];
     if (self) {
+        self.data   = data;
         self.cursor = (const char *)data.bytes;
-        self.end    = (const char *)data.bytes + data.length - 1;
+        self.end    = (const char *)data.bytes + data.length;
     }
     return self;
 }
@@ -78,41 +80,45 @@ - (UInt16)parseShortUInt {
 }
 
 - (char)parseOctet {
-    return *((self.cursor)++);
+    if (self.cursor + 1 > self.end) {
+        return 0;
+    } else {
+        return *((self.cursor)++);
+    }
 }
 
-- (NSString *)parseShortString {
-    UInt8 length = *(self.cursor);
-    self.cursor++;
-    NSString *string = [NSString stringWithFormat:@"%.*s", length, self.cursor];
-    self.cursor += length;
-
-    return string;
+- (BOOL)parseBoolean {
+    return [self parseOctet] != 0;
 }
 
-- (NSString *)parseLongString {
-    if (!self.cursor || self.cursor + 4 > self.end) {
-        // throw or something
-    }
+- (NSString *)parseShortString {
+    UInt8 length = *self.cursor;
+    const char *expectedStringEnd = self.cursor + sizeof(length) + length;
 
-    unsigned int length = CFSwapInt32BigToHost(*(UInt32 *)self.cursor);
-    self.cursor += sizeof(length);
+    if (expectedStringEnd > self.end) {
+        return @"";
+    } else {
+        self.cursor++;
+        NSString *string = [NSString stringWithFormat:@"%.*s", length, self.cursor];
+        self.cursor += length;
 
-    if (self.cursor + length > self.end) {
-        // TODO: What to do if length == 4GiB
+        return string;
     }
-    NSString *string= [NSString stringWithFormat:@"%.*s", length, self.cursor];
-    self.cursor += length;
-
-    return string;
 }
 
-- (BOOL)parseBoolean {
-    if (!self.cursor || self.cursor + 1 > self.end) {
-        // throw or something
-    }
+- (NSString *)parseLongString {
+    UInt32 length = CFSwapInt32BigToHost(*(UInt32 *)self.cursor);
+    const char *expectedStringEnd = self.cursor + sizeof(length) + length;
+
+    if (expectedStringEnd > self.end) {
+        return @"";
+    } else {
+        self.cursor += sizeof(length);
+        NSString *string = [NSString stringWithFormat:@"%.*s", length, self.cursor];
+        self.cursor += length;
 
-    return *((self.cursor)++) != 0;
+        return string;
+    }
 }
 
 - (NSData *)parseLength:(UInt32)length {

diff --git a/RMQClientTests/AMQParserTest.swift b/RMQClientTests/AMQParserTest.swift
@@ -2,6 +2,25 @@ import XCTest
 
 class AMQParserTest: XCTestCase {
 
+    func testOctet() {
+        let parser = AMQParser(data: "\u{2}".dataUsingEncoding(NSUTF8StringEncoding)!)
+        XCTAssertEqual(2, parser.parseOctet())
+
+        for _ in 1...1000 {
+            XCTAssertEqual(0, parser.parseOctet())
+        }
+    }
+
+    func testBoolean() {
+        let parser = AMQParser(data: "\u{1}\u{0}".dataUsingEncoding(NSUTF8StringEncoding)!)
+        XCTAssertTrue(parser.parseBoolean())
+        XCTAssertFalse(parser.parseBoolean())
+
+        for _ in 1...1000 {
+            XCTAssertFalse(parser.parseBoolean())
+        }
+    }
+
     func testShortString() {
         let s = "PRECONDITION_FAILED - inequivalent arg 'durable' for queue 'rmqclient.integration-tests.E0B5A093-6B2E-402C-84F3-E93B59DF807B-71865-0003F85C24C90FC6' in vhost '/': received 'false' but current is 'true'"
         let data = NSMutableData()
@@ -14,4 +33,42 @@ class AMQParserTest: XCTestCase {
         XCTAssertEqual(s, parser.parseShortString())
     }
 
+    func testShortStringWhenAlreadyRead() {
+        let parser = AMQParser(data: "\u{4}aaaa".dataUsingEncoding(NSUTF8StringEncoding)!)
+        XCTAssertEqual("aaaa", parser.parseShortString())
+        for _ in 1...1000 {
+            XCTAssertEqual("", parser.parseShortString())
+        }
+    }
+
+    func testShortStringWhenNotEnoughDataToReadAfterLongString() {
+        let parser = AMQParser(data: "\u{0}\u{0}\u{0}\u{4}AAAA\u{4}BBB".dataUsingEncoding(NSUTF8StringEncoding)!)
+        XCTAssertEqual("AAAA", parser.parseLongString())
+        XCTAssertEqual("", parser.parseShortString())
+    }
+
+    func testLongString() {
+        let parser = AMQParser(data: "\u{0}\u{0}\u{0}\u{4}AAAAbbbb".dataUsingEncoding(NSUTF8StringEncoding)!)
+        XCTAssertEqual("AAAA", parser.parseLongString())
+    }
+
+    func testLongStringWhenAlreadyRead() {
+        let parser = AMQParser(data: "\u{0}\u{0}\u{0}\u{4}AAAA".dataUsingEncoding(NSUTF8StringEncoding)!)
+        XCTAssertEqual("AAAA", parser.parseLongString())
+        for _ in 1...1000 {
+            XCTAssertEqual("", parser.parseLongString())
+        }
+    }
+
+    func testLongStringWhenNotEnoughDataToRead() {
+        let parser = AMQParser(data: "\u{0}\u{0}\u{0}\u{4}AAA".dataUsingEncoding(NSUTF8StringEncoding)!)
+        XCTAssertEqual("", parser.parseLongString())
+    }
+
+    func testLongStringWhenNotEnoughDataToReadAfterShortString() {
+        let parser = AMQParser(data: "\u{4}BBBB\u{0}\u{0}\u{0}\u{4}AAA".dataUsingEncoding(NSUTF8StringEncoding)!)
+        XCTAssertEqual("BBBB", parser.parseShortString())
+        XCTAssertEqual("", parser.parseLongString())
+    }
+
 }