Fix malformed PUBLISH parsing and handle protocol errors safely

CocoaMQTTTests/CocoaMQTTReaderProtocolErrorTests.swift

@@ -0,0 +1,60 @@
+import Foundation
+import XCTest
+@testable import CocoaMQTT
+
+final class CocoaMQTTReaderProtocolErrorTests: XCTestCase {
+
+    private final class SocketSpy: CocoaMQTTSocketProtocol {
+        var enableSSL: Bool = false
+        private(set) var disconnectCount = 0
+
+        func setDelegate(_ theDelegate: CocoaMQTTSocketDelegate?, delegateQueue: DispatchQueue?) {}
+        func connect(toHost host: String, onPort port: UInt16) throws {}
+        func connect(toHost host: String, onPort port: UInt16, withTimeout timeout: TimeInterval) throws {}
+        func disconnect() { disconnectCount += 1 }
+        func readData(toLength length: UInt, withTimeout timeout: TimeInterval, tag: Int) {}
+        func write(_ data: Data, withTimeout timeout: TimeInterval, tag: Int) {}
+    }
+
+    private final class ReaderDelegateSpy: CocoaMQTTReaderDelegate {
+        private(set) var publishCount = 0
+
+        func didReceive(_ reader: CocoaMQTTReader, connack: FrameConnAck) {}
+        func didReceive(_ reader: CocoaMQTTReader, publish: FramePublish) { publishCount += 1 }
+        func didReceive(_ reader: CocoaMQTTReader, puback: FramePubAck) {}
+        func didReceive(_ reader: CocoaMQTTReader, pubrec: FramePubRec) {}
+        func didReceive(_ reader: CocoaMQTTReader, pubrel: FramePubRel) {}
+        func didReceive(_ reader: CocoaMQTTReader, pubcomp: FramePubComp) {}
+        func didReceive(_ reader: CocoaMQTTReader, suback: FrameSubAck) {}
+        func didReceive(_ reader: CocoaMQTTReader, unsuback: FrameUnsubAck) {}
+        func didReceive(_ reader: CocoaMQTTReader, pingresp: FramePingResp) {}
+    }
+
+    func testMalformedPublishDisconnectsSocket() {
+        CocoaMQTTStorage()?.setMQTTVersion("3.1.1")
+
+        let socket = SocketSpy()
+        let delegate = ReaderDelegateSpy()
+        let reader = CocoaMQTTReader(socket: socket, delegate: delegate)
+
+        reader.headerReady(FrameType.publish.rawValue)
+        reader.lengthReady(0x06)
+        reader.payloadReady(Data([0x00, 0x00, 0x41, 0x41, 0x41, 0x41]))
+
+        XCTAssertEqual(socket.disconnectCount, 1)
+        XCTAssertEqual(delegate.publishCount, 0)
+    }
+
+    func testUnknownFrameTypeDisconnectsSocket() {
+        let socket = SocketSpy()
+        let delegate = ReaderDelegateSpy()
+        let reader = CocoaMQTTReader(socket: socket, delegate: delegate)
+
+        reader.headerReady(0x00)
+        reader.lengthReady(0x00)
+
+        XCTAssertEqual(socket.disconnectCount, 1)
+        XCTAssertEqual(delegate.publishCount, 0)
+    }
+}
+

CocoaMQTTTests/FrameTests.swift

@@ -142,6 +142,28 @@ class FrameTests: XCTestCase {
         XCTAssertEqual(f1?.payload().count, 0)
     }
 
+    func testFramePublishRejectsZeroLengthTopic() {
+        CocoaMQTTStorage()?.setMQTTVersion("3.1.1")
+        let frame = FramePublish(packetFixedHeaderType: FrameType.publish.rawValue, bytes: [0x00, 0x00, 0x41, 0x41, 0x41, 0x41])
+        XCTAssertNil(frame)
+    }
+
+    func testFramePublishRejectsZeroLengthTopicInMQTT5WithoutAlias() {
+        CocoaMQTTStorage()?.setMQTTVersion("5.0")
+        defer { CocoaMQTTStorage()?.setMQTTVersion("3.1.1") }
+
+        let frame = FramePublish(packetFixedHeaderType: FrameType.publish.rawValue, bytes: [0x00, 0x00, 0x00])
+        XCTAssertNil(frame)
+    }
+
+    func testFramePublishAllowsZeroLengthTopicInMQTT5WithAlias() {
+        CocoaMQTTStorage()?.setMQTTVersion("5.0")
+        defer { CocoaMQTTStorage()?.setMQTTVersion("3.1.1") }
+
+        let frame = FramePublish(packetFixedHeaderType: FrameType.publish.rawValue, bytes: [0x00, 0x00, 0x03, 0x23, 0x00, 0x01])
+        XCTAssertEqual(frame?.topic, "")
+    }
+
     func testFramePubAck() {
 
         var puback = FramePubAck(msgid: 0x1010)

Source/CocoaMQTTReader.swift

@@ -109,8 +109,7 @@ class CocoaMQTTReader {
     private func frameReady() {
 
         guard let frameType = FrameType(rawValue: UInt8(header & 0xF0)) else {
-            printError("Received unknown frame type, header: \(header), data:\(data)")
-            readHeader()
+            protocolError("Received unknown frame type, header: \(header), data:\(data)")
             return
         }
 
@@ -119,65 +118,71 @@ class CocoaMQTTReader {
         switch frameType {
         case .connack:
             guard let connack = FrameConnAck(packetFixedHeaderType: header, bytes: data) else {
-                printError("Reader parse \(frameType) failed, data: \(data)")
-                break
+                protocolError("Reader parse \(frameType) failed, data: \(data)")
+                return
             }
             delegate?.didReceive(self, connack: connack)
         case .publish:
             guard let publish = FramePublish(packetFixedHeaderType: header, bytes: data) else {
-                printError("Reader parse \(frameType) failed, data: \(data)")
-                break
+                protocolError("Reader parse \(frameType) failed, data: \(data)")
+                return
             }
             delegate?.didReceive(self, publish: publish)
         case .puback:
             guard let puback = FramePubAck(packetFixedHeaderType: header, bytes: data) else {
-                printError("Reader parse \(frameType) failed, data: \(data)")
-                break
+                protocolError("Reader parse \(frameType) failed, data: \(data)")
+                return
             }
             delegate?.didReceive(self, puback: puback)
         case .pubrec:
             guard let pubrec = FramePubRec(packetFixedHeaderType: header, bytes: data) else {
-                printError("Reader parse \(frameType) failed, data: \(data)")
-                break
+                protocolError("Reader parse \(frameType) failed, data: \(data)")
+                return
             }
             delegate?.didReceive(self, pubrec: pubrec)
         case .pubrel:
             guard let pubrel = FramePubRel(packetFixedHeaderType: header, bytes: data) else {
-                printError("Reader parse \(frameType) failed, data: \(data)")
-                break
+                protocolError("Reader parse \(frameType) failed, data: \(data)")
+                return
             }
             delegate?.didReceive(self, pubrel: pubrel)
         case .pubcomp:
             guard let pubcomp = FramePubComp(packetFixedHeaderType: header, bytes: data) else {
-                printError("Reader parse \(frameType) failed, data: \(data)")
-                break
+                protocolError("Reader parse \(frameType) failed, data: \(data)")
+                return
             }
             delegate?.didReceive(self, pubcomp: pubcomp)
         case .suback:
             guard let frame = FrameSubAck(packetFixedHeaderType: header, bytes: data) else {
-                printError("Reader parse \(frameType) failed, data: \(data)")
-                break
+                protocolError("Reader parse \(frameType) failed, data: \(data)")
+                return
             }
             delegate?.didReceive(self, suback: frame)
         case .unsuback:
             guard let frame = FrameUnsubAck(packetFixedHeaderType: header, bytes: data) else {
-                printError("Reader parse \(frameType) failed, data: \(data)")
-                break
+                protocolError("Reader parse \(frameType) failed, data: \(data)")
+                return
             }
             delegate?.didReceive(self, unsuback: frame)
         case .pingresp:
             guard let frame = FramePingResp(packetFixedHeaderType: header, bytes: data) else {
-                printError("Reader parse \(frameType) failed, data: \(data)")
-                break
+                protocolError("Reader parse \(frameType) failed, data: \(data)")
+                return
             }
             delegate?.didReceive(self, pingresp: frame)
         default:
-            break
+            protocolError("Received unsupported frame type \(frameType), data: \(data)")
+            return
         }
 
         readHeader()
     }
 
+    private func protocolError(_ reason: String) {
+        printError(reason)
+        socket.disconnect()
+    }
+
     private func reset() {
         length = 0
         multiply = 1

Source/FramePublish.swift

@@ -156,14 +156,14 @@ extension FramePublish: InitialWithBytes {
             return nil
         }
 
-        let len = UInt16(bytes[0]) << 8 + UInt16(bytes[1])
+        let topicLength = Int(UInt16(bytes[0]) << 8 | UInt16(bytes[1]))
+        let topicStart = 2
+        let topicEnd = topicStart + topicLength
 
-        // 2 is packetFixedHeaderType length
-        var pos = 2 + Int(len)
-
-        if bytes.count < pos {
+        if bytes.count < topicEnd {
             return nil
         }
+        var pos = topicEnd
 
         // msgid
         if (packetFixedHeaderType & 0x06) >> 1 == CocoaMQTTQoS.qos0.rawValue {
@@ -189,6 +189,14 @@ extension FramePublish: InitialWithBytes {
             if data.propertyLength != 0 {
                 pos += data.propertyLength!
             }
+            if pos > bytes.count {
+                return nil
+            }
+
+            // MQTT 5.0: Topic Name may be empty only when Topic Alias is present.
+            if data.topic.isEmpty && data.topicAlias == nil {
+                return nil
+            }
 
             // MQTT 5.0
             self.mqtt5Topic = data.topic
@@ -201,9 +209,13 @@ extension FramePublish: InitialWithBytes {
 
         } else {
             // MQTT 3.1.1
-            if let data = NSString(bytes: [UInt8](bytes[2...(pos-1)]), length: Int(len), encoding: String.Encoding.utf8.rawValue) {
-                topic =  data as String
+            guard topicLength > 0 else {
+                return nil
+            }
+            guard let recTopic = String(bytes: bytes[topicStart..<topicEnd], encoding: .utf8) else {
+                return nil
             }
+            topic = recTopic
         }
 
         // payload