Don’t crash sourcekit-lsp if a known message is missing a field

Previously, we `fatalError`ing when `JSONDecoder` failed to decode a message from the client. Instead of crashing, try recovering from such invalid messages as best as possible. If we know that the state might  have gotten out of sync with the client, show a notification message to the user, asking them to file an issue.

rdar://112991102

Author: ahoppen

Sources/LSPTestSupport/TestJSONRPCConnection.swift

@@ -197,7 +197,7 @@ public final class TestServer: MessageHandler {
 
 private let testMessageRegistry = MessageRegistry(
   requests: [EchoRequest.self, EchoError.self],
-  notifications: [EchoNotification.self]
+  notifications: [EchoNotification.self, ShowMessageNotification.self]
 )
 
 #if compiler(<5.11)

Sources/LanguageServerProtocolJSONRPC/JSONRPCConnection.swift

@@ -56,7 +56,7 @@ public final class JSONRPCConnection: Connection {
   ///  - `init`: Reference to `JSONRPCConnection` trivially can't have escaped to other isolation domains yet.
   ///  - `start`: Is required to be called in the same serial region as the initializer, so
   ///    `JSONRPCConnection` can't have escaped to other isolation domains yet.
-  ///  - `_close`: Synchronized on `queue`.
+  ///  - `closeAssumingOnQueue`: Synchronized on `queue`.
   ///  - `readyToSend`: Synchronized on `queue`.
   ///  - `deinit`: Can also only trivially be called once.
   private nonisolated(unsafe) var state: State
@@ -230,6 +230,131 @@ public final class JSONRPCConnection: Connection {
     }
   }
 
+  /// Send a notification to the client that informs the user about a message decoding error and tells them to file an
+  /// issue.
+  ///
+  /// `message` describes what has gone wrong to the user.
+  ///
+  /// - Important: Must be called on `queue`
+  private func sendMessageDecodingErrorNotificationToClient(message: String) {
+    dispatchPrecondition(condition: .onQueue(queue))
+    let showMessage = ShowMessageNotification(
+      type: .error,
+      message: """
+        \(message). Please run 'sourcekit-lsp diagnose' to file an issue.
+        """
+    )
+    self.send(.notification(showMessage))
+  }
+
+  /// Decode a single JSONRPC message from the given `messageBytes`.
+  ///
+  /// `messageBytes` should be valid JSON, ie. this is the message sent from the client without the `Content-Length`
+  /// header.
+  ///
+  /// If an error occurs during message parsing, this tries to recover as gracefully as possible and returns `nil`.
+  /// Callers should consider the message handled and ignore it when this function returns `nil`.
+  ///
+  /// - Important: Must be called on `queue`
+  private func decodeJSONRPCMessage(messageBytes: Slice<UnsafeBufferPointer<UInt8>>) -> JSONRPCMessage? {
+    dispatchPrecondition(condition: .onQueue(queue))
+    let decoder = JSONDecoder()
+
+    // Set message registry to use for model decoding.
+    decoder.userInfo[.messageRegistryKey] = messageRegistry
+
+    // Setup callback for response type.
+    decoder.userInfo[.responseTypeCallbackKey] = { (id: RequestID) -> ResponseType.Type? in
+      guard let outstanding = self.outstandingRequests[id] else {
+        logger.error("Unknown request for \(id, privacy: .public)")
+        return nil
+      }
+      return outstanding.responseType
+    }
+
+    do {
+      let pointer = UnsafeMutableRawPointer(mutating: UnsafeBufferPointer(rebasing: messageBytes).baseAddress!)
+      return try decoder.decode(
+        JSONRPCMessage.self,
+        from: Data(bytesNoCopy: pointer, count: messageBytes.count, deallocator: .none)
+      )
+    } catch let error as MessageDecodingError {
+      logger.fault("Failed to decode message: \(error.forLogging)")
+      logger.fault("Malformed message: \(String(bytes: messageBytes, encoding: .utf8) ?? "<invalid UTF-8>")")
+
+      // We failed to decode the message. Under those circumstances try to behave as LSP-conforming as possible.
+      // Always log at the fault level so that we know something is going wrong from the logs.
+      //
+      // The pattern below is to handle the message in the best possible way and then `return nil` to acknowledge the
+      // handling. That way the compiler enforces that we handle all code paths.
+      switch error.messageKind {
+      case .request:
+        if let id = error.id {
+          // If we know it was a request and we have the request ID, simply reply to the request and tell the client
+          // that we couldn't parse it. That complies with LSP that all requests should eventually get a response.
+          logger.fault(
+            "Replying to request \(id, privacy: .public) with error response because we failed to decode the request"
+          )
+          self.send(.errorResponse(ResponseError(error), id: id))
+          return nil
+        }
+        // If we don't know the ID of the request, ignore it and show a notification to the user.
+        // That way the user at least knows that something is going wrong even if the client never gets a response
+        // for the request.
+        logger.fault("Ignoring request because we failed to decode the request and don't have a request ID")
+        sendMessageDecodingErrorNotificationToClient(message: "sourcekit-lsp failed to decode a request")
+        return nil
+      case .response:
+        if let id = error.id {
+          if let outstanding = self.outstandingRequests.removeValue(forKey: id) {
+            // If we received a response to a request we sent to the client, assume that the client responded with an
+            // error. That complies with LSP that all requests should eventually get a response.
+            logger.fault(
+              "Assuming an error response to request \(id, privacy: .public) because response from client could not be decoded"
+            )
+            outstanding.replyHandler(.failure(ResponseError(error)))
+            return nil
+          }
+          // If there's an error in the response but we don't even know about the request, we can ignore it.
+          logger.fault(
+            "Ignoring response to request \(id, privacy: .public) because it could not be decoded and given request ID is unknown"
+          )
+          return nil
+        }
+        // And if we can't even recover the ID the response is for, we drop it. This means that whichever code in
+        // sourcekit-lsp sent the request will probably never get a reply but there's nothing we can do about that.
+        // Ideally requests sent from sourcekit-lsp to the client would have some kind of timeout anyway.
+        logger.fault("Ignoring response because its request ID could not be recovered")
+        return nil
+      case .notification:
+        if error.code == .methodNotFound {
+          // If we receive a notification we don't know about, this might be a client sending a new LSP notification
+          // that we don't know about. It can't be very critical so we ignore it without bothering the user with an
+          // error notification.
+          logger.fault("Ignoring notification because we don't know about it's method")
+          return nil
+        }
+        // Ignoring any other notification might result in corrupted behavior. For example, ignoring a
+        // `textDocument/didChange` will result in an out-of-sync state between the editor and sourcekit-lsp.
+        // Warn the user about the error.
+        logger.fault("Ignoring notification that may cause corrupted behavior")
+        sendMessageDecodingErrorNotificationToClient(message: "sourcekit-lsp failed to decode a notification")
+        return nil
+      case .unknown:
+        // We don't know what has gone wrong. This could be any level of badness. Inform the user about it.
+        logger.fault("Ignoring unknown message")
+        sendMessageDecodingErrorNotificationToClient(message: "sourcekit-lsp failed to decode a message")
+        return nil
+      }
+    } catch {
+      // We don't know what has gone wrong. This could be any level of badness. Inform the user about it and ignore the
+      // message.
+      logger.fault("ignoring unknown message")
+      sendMessageDecodingErrorNotificationToClient(message: "sourcekit-lsp failed to decode an unknown message")
+      return nil
+    }
+  }
+
   /// Whether we can send messages in the current state.
   ///
   /// - parameter shouldLog: Whether to log an info message if not ready.
@@ -250,69 +375,30 @@ public final class JSONRPCConnection: Connection {
   /// - Important: Must be called on `queue`
   func parseAndHandleMessages(from bytes: UnsafeBufferPointer<UInt8>) -> UnsafeBufferPointer<UInt8>.SubSequence {
     dispatchPrecondition(condition: .onQueue(queue))
-    let decoder = JSONDecoder()
-
-    // Set message registry to use for model decoding.
-    decoder.userInfo[.messageRegistryKey] = messageRegistry
-
-    // Setup callback for response type.
-    decoder.userInfo[.responseTypeCallbackKey] =
-      { id in
-        guard let outstanding = self.outstandingRequests[id] else {
-          logger.error("Unknown request for \(id, privacy: .public)")
-          return nil
-        }
-        return outstanding.responseType
-      } as JSONRPCMessage.ResponseTypeCallback
 
     var bytes = bytes[...]
 
     MESSAGE_LOOP: while true {
+      // Split the messages based on the Content-Length header.
+      let messageBytes: Slice<UnsafeBufferPointer<UInt8>>
       do {
-        guard let ((messageBytes, _), rest) = try bytes.jsonrpcSplitMessage() else {
+        guard let (header: _, message: message, rest: rest) = try bytes.jsonrpcSplitMessage() else {
           return bytes
         }
+        messageBytes = message
         bytes = rest
-
-        let pointer = UnsafeMutableRawPointer(mutating: UnsafeBufferPointer(rebasing: messageBytes).baseAddress!)
-        let message = try decoder.decode(
-          JSONRPCMessage.self,
-          from: Data(bytesNoCopy: pointer, count: messageBytes.count, deallocator: .none)
-        )
-
-        handle(message)
-      } catch let error as MessageDecodingError {
-        switch error.messageKind {
-        case .request:
-          if let id = error.id {
-            queue.async {
-              self.send(.errorResponse(ResponseError(error), id: id))
-            }
-            continue MESSAGE_LOOP
-          }
-        case .response:
-          if let id = error.id {
-            if let outstanding = self.outstandingRequests.removeValue(forKey: id) {
-              outstanding.replyHandler(.failure(ResponseError(error)))
-            } else {
-              logger.error("error in response to unknown request \(id, privacy: .public) \(error.forLogging)")
-            }
-            continue MESSAGE_LOOP
-          }
-        case .notification:
-          if error.code == .methodNotFound {
-            logger.error("ignoring unknown notification \(error.forLogging)")
-            continue MESSAGE_LOOP
-          }
-        case .unknown:
-          break
-        }
-        // FIXME: graceful shutdown?
-        fatalError("fatal error encountered decoding message \(error)")
       } catch {
-        // FIXME: graceful shutdown?
-        fatalError("fatal error encountered decoding message \(error)")
+        // We failed to parse the message header. There isn't really much we can do to recover because we lost our
+        // anchor in the stream where new messages start. Crashing and letting ourselves be restarted by the client is
+        // probably the best option.
+        sendMessageDecodingErrorNotificationToClient(message: "Failed to find next message in connection to editor")
+        fatalError("fatal error encountered while splitting JSON RPC messages \(error)")
+      }
+
+      guard let message = decodeJSONRPCMessage(messageBytes: messageBytes) else {
+        continue
       }
+      handle(message)
     }
   }
 

Sources/LanguageServerProtocolJSONRPC/MessageSplitting.swift

@@ -15,7 +15,7 @@ import LanguageServerProtocol
 public struct JSONRPCMessageHeader: Hashable {
   static let contentLengthKey: [UInt8] = [UInt8]("Content-Length".utf8)
   static let separator: [UInt8] = [UInt8]("\r\n".utf8)
-  static let colon: UInt8 = ":".utf8.first!
+  static let colon: UInt8 = UInt8(ascii: ":")
   static let invalidKeyBytes: [UInt8] = [colon] + separator
 
   public var contentLength: Int? = nil
@@ -25,21 +25,29 @@ public struct JSONRPCMessageHeader: Hashable {
   }
 }
 
-extension RandomAccessCollection where Element == UInt8 {
-
-  /// Returns the first message range and header in `self`, or nil.
-  public func jsonrpcSplitMessage()
-    throws -> ((SubSequence, header: JSONRPCMessageHeader), SubSequence)?
-  {
+extension RandomAccessCollection<UInt8> {
+  /// Tries to parse a single message from this collection of bytes.
+  ///
+  /// If an entire message could be found, returns
+  ///  - header (representing `Content-Length:<length>\r\n\r\n`)
+  ///  - message: The data that represents the actual message as JSON
+  ///  - rest: The remaining bytes that haven't weren't part of the first message in this collection
+  ///
+  /// If a `Content-Length` header could be found but the collection doesn't have enough bytes for the entire message
+  /// (eg. because the `Content-Length` header has been transmitted yet but not the entire message), returns `nil`.
+  /// Callers should call this method again once more data is available.
+  @_spi(Testing)
+  public func jsonrpcSplitMessage() throws -> (header: JSONRPCMessageHeader, message: SubSequence, rest: SubSequence)? {
     guard let (header, rest) = try jsonrcpParseHeader() else { return nil }
     guard let contentLength = header.contentLength else {
       throw MessageDecodingError.parseError("missing Content-Length header")
     }
     if contentLength > rest.count { return nil }
-    return ((rest.prefix(contentLength), header: header), rest.dropFirst(contentLength))
+    return (header: header, message: rest.prefix(contentLength), rest: rest.dropFirst(contentLength))
   }
 
-  public func jsonrcpParseHeader() throws -> (JSONRPCMessageHeader, SubSequence)? {
+  @_spi(Testing)
+  public func jsonrcpParseHeader() throws -> (header: JSONRPCMessageHeader, rest: SubSequence)? {
     var header = JSONRPCMessageHeader()
     var slice = self[...]
     while let (kv, rest) = try slice.jsonrpcParseHeaderField() {
@@ -62,6 +70,7 @@ extension RandomAccessCollection where Element == UInt8 {
     return nil
   }
 
+  @_spi(Testing)
   public func jsonrpcParseHeaderField() throws -> ((key: SubSequence, value: SubSequence)?, SubSequence)? {
     if starts(with: JSONRPCMessageHeader.separator) {
       return (nil, dropFirst(JSONRPCMessageHeader.separator.count))
@@ -85,11 +94,9 @@ extension RandomAccessCollection where Element == UInt8 {
 }
 
 extension RandomAccessCollection where Element: Equatable {
-
   /// Returns the first index where the specified subsequence appears or nil.
   @inlinable
   public func firstIndex(of pattern: some RandomAccessCollection<Element>) -> Index? {
-
     if pattern.isEmpty {
       return startIndex
     }

Tests/LanguageServerProtocolJSONRPCTests/ConnectionTests.swift

@@ -279,4 +279,32 @@ class ConnectionTests: XCTestCase {
       }
     }
   }
+
+  func testMessageWithMissingParameter() async throws {
+    let expectation = self.expectation(description: "Received ShowMessageNotification")
+    await connection.client.appendOneShotNotificationHandler { (note: ShowMessageNotification) in
+      XCTAssertEqual(note.type, .error)
+      expectation.fulfill()
+    }
+
+    let messageContents = """
+      {
+        "method": "test_server/echo_note",
+        "jsonrpc": "2.0",
+        "params": {}
+      }
+      """
+    connection.clientToServerConnection.send(message: messageContents)
+
+    try await self.fulfillmentOfOrThrow([expectation])
+  }
+}
+
+fileprivate extension JSONRPCConnection {
+  func send(message: String) {
+    let messageWithHeader = "Content-Length: \(message.utf8.count)\r\n\r\n\(message)".data(using: .utf8)!
+    messageWithHeader.withUnsafeBytes { bytes in
+      send(_rawData: DispatchData(bytes: bytes))
+    }
+  }
 }

Tests/LanguageServerProtocolJSONRPCTests/MessageParsingTests.swift

@@ -11,7 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 import LanguageServerProtocol
-import LanguageServerProtocolJSONRPC
+@_spi(Testing) import LanguageServerProtocolJSONRPC
 import XCTest
 
 final class MessageParsingTests: XCTestCase {
@@ -25,7 +25,7 @@ final class MessageParsingTests: XCTestCase {
       line: UInt = #line
     ) throws {
       let bytes: [UInt8] = [UInt8](string.utf8)
-      guard let ((content, header), rest) = try bytes.jsonrpcSplitMessage() else {
+      guard let (header, content, rest) = try bytes.jsonrpcSplitMessage() else {
         XCTAssert(restLen == nil, "expected non-empty field", file: file, line: line)
         return
       }