Merge pull request #659 from emqx/fix/publish-protocol-error-handling

Fix malformed PUBLISH parsing and handle protocol errors safely

Author: JianBo He

CocoaMQTTTests/CocoaMQTTReaderProtocolErrorTests.swift

@@ -0,0 +1,105 @@
+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
+        private(set) var disconnectCount = 0
+        private(set) var authCount = 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 didReceive(_ reader: CocoaMQTTReader, disconnect: FrameDisconnect) { disconnectCount += 1 }
+        func didReceive(_ reader: CocoaMQTTReader, auth: FrameAuth) { authCount += 1 }
+    }
+
+    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)
+    }
+
+    func testMQTT5DisconnectFrameDoesNotProtocolError() {
+        CocoaMQTTStorage()?.setMQTTVersion("5.0")
+
+        let socket = SocketSpy()
+        let delegate = ReaderDelegateSpy()
+        let reader = CocoaMQTTReader(socket: socket, delegate: delegate)
+
+        reader.headerReady(FrameType.disconnect.rawValue)
+        reader.lengthReady(0x00)
+
+        XCTAssertEqual(socket.disconnectCount, 0)
+        XCTAssertEqual(delegate.disconnectCount, 1)
+    }
+
+    func testMQTT5AuthFrameDoesNotProtocolError() {
+        CocoaMQTTStorage()?.setMQTTVersion("5.0")
+
+        let socket = SocketSpy()
+        let delegate = ReaderDelegateSpy()
+        let reader = CocoaMQTTReader(socket: socket, delegate: delegate)
+
+        reader.headerReady(FrameType.auth.rawValue)
+        reader.lengthReady(0x00)
+
+        XCTAssertEqual(socket.disconnectCount, 0)
+        XCTAssertEqual(delegate.authCount, 1)
+    }
+
+    func testMQTT311RejectsMQTT5OnlyDisconnectFrame() {
+        CocoaMQTTStorage()?.setMQTTVersion("3.1.1")
+
+        let socket = SocketSpy()
+        let delegate = ReaderDelegateSpy()
+        let reader = CocoaMQTTReader(socket: socket, delegate: delegate)
+
+        reader.headerReady(FrameType.disconnect.rawValue)
+        reader.lengthReady(0x00)
+
+        XCTAssertEqual(socket.disconnectCount, 1)
+        XCTAssertEqual(delegate.disconnectCount, 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/CocoaMQTT.swift

@@ -794,4 +794,14 @@ extension CocoaMQTT: CocoaMQTTReaderDelegate {
         delegate?.mqttDidReceivePong(self)
         didReceivePong(self)
     }
+
+    func didReceive(_ reader: CocoaMQTTReader, disconnect: FrameDisconnect) {
+        printWarning("Received DISCONNECT in MQTT 3.1.1 mode, closing socket")
+        internal_disconnect()
+    }
+
+    func didReceive(_ reader: CocoaMQTTReader, auth: FrameAuth) {
+        printWarning("Received AUTH in MQTT 3.1.1 mode, closing socket")
+        internal_disconnect()
+    }
 }

Source/CocoaMQTT5.swift

@@ -702,13 +702,15 @@ extension CocoaMQTT5: CocoaMQTTSocketDelegate {
 extension CocoaMQTT5: CocoaMQTTReaderDelegate {
 
     func didReceive(_ reader: CocoaMQTTReader, disconnect: FrameDisconnect) {
-        delegate?.mqtt5(self, didReceiveDisconnectReasonCode: disconnect.receiveReasonCode!)
-        didDisconnectReasonCode(self, disconnect.receiveReasonCode!)
+        let reasonCode = disconnect.receiveReasonCode ?? .normalDisconnection
+        delegate?.mqtt5(self, didReceiveDisconnectReasonCode: reasonCode)
+        didDisconnectReasonCode(self, reasonCode)
     }
 
     func didReceive(_ reader: CocoaMQTTReader, auth: FrameAuth) {
-        delegate?.mqtt5(self, didReceiveAuthReasonCode: auth.receiveReasonCode!)
-        didAuthReasonCode(self, auth.receiveReasonCode!)
+        let reasonCode = auth.receiveReasonCode ?? .success
+        delegate?.mqtt5(self, didReceiveAuthReasonCode: reasonCode)
+        didAuthReasonCode(self, reasonCode)
     }
 
     func didReceive(_ reader: CocoaMQTTReader, connack: FrameConnAck) {

Source/CocoaMQTTReader.swift

@@ -35,6 +35,10 @@ protocol CocoaMQTTReaderDelegate: AnyObject {
     func didReceive(_ reader: CocoaMQTTReader, unsuback: FrameUnsubAck)
 
     func didReceive(_ reader: CocoaMQTTReader, pingresp: FramePingResp)
+
+    func didReceive(_ reader: CocoaMQTTReader, disconnect: FrameDisconnect)
+
+    func didReceive(_ reader: CocoaMQTTReader, auth: FrameAuth)
 }
 
 class CocoaMQTTReader {
@@ -109,8 +113,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 +122,95 @@ 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)
+        case .disconnect:
+            guard isMQTT5ProtocolVersion() else {
+                protocolError("Reader received MQTT5-only frame \(frameType) in non-MQTT5 mode, data: \(data)")
+                return
+            }
+            guard let frame = FrameDisconnect(packetFixedHeaderType: header, bytes: data) else {
+                protocolError("Reader parse \(frameType) failed, data: \(data)")
+                return
+            }
+            delegate?.didReceive(self, disconnect: frame)
+        case .auth:
+            guard isMQTT5ProtocolVersion() else {
+                protocolError("Reader received MQTT5-only frame \(frameType) in non-MQTT5 mode, data: \(data)")
+                return
+            }
+            guard let frame = FrameAuth(packetFixedHeaderType: header, bytes: data) else {
+                protocolError("Reader parse \(frameType) failed, data: \(data)")
+                return
+            }
+            delegate?.didReceive(self, auth: frame)
         default:
-            break
+            protocolError("Received unsupported frame type \(frameType), data: \(data)")
+            return
         }
 
         readHeader()
     }
 
+    private func protocolError(_ reason: String) {
+        printError(reason)
+        socket.disconnect()
+    }
+
+    private func isMQTT5ProtocolVersion() -> Bool {
+        return CocoaMQTTStorage()?.queryMQTTVersion() == "5.0"
+    }
+
     private func reset() {
         length = 0
         multiply = 1

Source/FrameAuth.swift

@@ -70,8 +70,18 @@ extension FrameAuth {
 extension FrameAuth: InitialWithBytes {
 
     init?(packetFixedHeaderType: UInt8, bytes: [UInt8]) {
-
-        receiveReasonCode = CocoaMQTTAUTHReasonCode(rawValue: bytes[0])
+        var protocolVersion = ""
+        if let storage = CocoaMQTTStorage() {
+            protocolVersion = storage.queryMQTTVersion()
+        }
+        guard protocolVersion == "5.0" else {
+            return nil
+        }
+        if bytes.isEmpty {
+            receiveReasonCode = .success
+        } else {
+            receiveReasonCode = CocoaMQTTAUTHReasonCode(rawValue: bytes[0])
+        }
     }
 
 }

Source/FrameDisconnect.swift

@@ -109,9 +109,13 @@ extension FrameDisconnect: InitialWithBytes {
         }
 
         if protocolVersion == "5.0" {
-            if bytes.count > 0 {
+            if bytes.isEmpty {
+                receiveReasonCode = .normalDisconnection
+            } else {
                 receiveReasonCode = CocoaMQTTDISCONNECTReasonCode(rawValue: bytes[0])
             }
+        } else {
+            return nil
         }
     }
 

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