Merge pull request #2179 from hamishknight/bound-check

Avoid crashing on invalid range in `editDocument`

Author: hamishknight

Sources/SKUtilities/LineTable.swift

@@ -115,45 +115,47 @@ extension LineTable {
     self.replace(fromLine: fromLine, utf16Offset: fromOff, toLine: toLine, utf16Offset: toOff, with: replacement)
   }
 
-  /// Replace the line table's `content` in the given range and update the line data.
-  ///
-  /// - parameter fromLine: Starting line number (zero-based).
-  /// - parameter fromOff: Starting UTF-8 column offset (zero-based).
-  /// - parameter toLine: Ending line number (zero-based).
-  /// - parameter toOff: Ending UTF-8 column offset (zero-based).
-  /// - parameter replacement: The new text for the given range.
-  @inlinable
-  mutating package func replace(
-    fromLine: Int,
-    utf8Offset fromOff: Int,
-    toLine: Int,
-    utf8Offset toOff: Int,
-    with replacement: String
-  ) {
-    let start = self.stringIndexOf(line: fromLine, utf8Column: fromOff)
-    let end = self.stringIndexOf(line: toLine, utf8Column: toOff)
-
-    var newText = self.content
-    newText.replaceSubrange(start..<end, with: replacement)
+  private struct OutOfBoundsError: Error, CustomLogStringConvertible {
+    // Note we use tuples here rather than Range since the latter would assert
+    // that upperBound >= lowerBound.
+    var utf8Range: (lower: Int, upper: Int)
+    var utf8Bounds: (lower: Int, upper: Int)
+
+    var description: String {
+      """
+      \(utf8Range.lower)..<\(utf8Range.upper) is out of bounds \
+      \(utf8Bounds.lower)..<\(utf8Bounds.upper)
+      """
+    }
 
-    self = LineTable(newText)
+    var redactedDescription: String {
+      description
+    }
   }
 
   /// Replace the line table's `content` in the given range and update the line data.
+  /// If the given range is out-of-bounds, throws an error.
   ///
-  /// - parameter fromLine: Starting line number (zero-based).
-  /// - parameter fromOff: Starting UTF-8 column offset (zero-based).
-  /// - parameter toLine: Ending line number (zero-based).
-  /// - parameter toOff: Ending UTF-8 column offset (zero-based).
+  /// - parameter utf8Offset: Starting UTF-8 offset (zero-based).
+  /// - parameter length: UTF-8 length.
   /// - parameter replacement: The new text for the given range.
   @inlinable
-  mutating package func replace(
+  mutating package func tryReplace(
     utf8Offset fromOff: Int,
     length: Int,
     with replacement: String
-  ) {
-    let start = content.utf8.index(content.startIndex, offsetBy: fromOff)
-    let end = content.utf8.index(content.startIndex, offsetBy: fromOff + length)
+  ) throws {
+    let utf8 = self.content.utf8
+    guard
+      fromOff >= 0, length >= 0,
+      let start = utf8.index(utf8.startIndex, offsetBy: fromOff, limitedBy: utf8.endIndex),
+      let end = utf8.index(start, offsetBy: length, limitedBy: utf8.endIndex)
+    else {
+      throw OutOfBoundsError(
+        utf8Range: (lower: fromOff, upper: fromOff + length),
+        utf8Bounds: (lower: 0, upper: utf8.count)
+      )
+    }
 
     var newText = self.content
     newText.replaceSubrange(start..<end, with: replacement)

Sources/SwiftSourceKitPlugin/CodeCompletion/Connection.swift

@@ -140,25 +140,11 @@ final class Connection {
       return
     }
 
-    document.lineTable.replace(utf8Offset: offset, length: length, with: newText)
-
-    sourcekitd.ideApi.set_file_contents(impl, path, document.lineTable.content)
-  }
-
-  func editDocument(path: String, edit: TextEdit) {
-    guard let document = documents[path] else {
-      logger.error("Document at '\(path)' is not open")
-      return
+    // Try replace the range, ignoring an invalid input. This matches SourceKit's
+    // behavior.
+    orLog("Replacing text") {
+      try document.lineTable.tryReplace(utf8Offset: offset, length: length, with: newText)
     }
-
-    document.lineTable.replace(
-      fromLine: edit.range.lowerBound.line - 1,
-      utf8Offset: edit.range.lowerBound.utf8Column - 1,
-      toLine: edit.range.upperBound.line - 1,
-      utf8Offset: edit.range.upperBound.utf8Column - 1,
-      with: edit.newText
-    )
-
     sourcekitd.ideApi.set_file_contents(impl, path, document.lineTable.content)
   }
 

Tests/SwiftSourceKitPluginTests/SwiftSourceKitPluginTests.swift

@@ -405,6 +405,80 @@ final class SwiftSourceKitPluginTests: XCTestCase {
     XCTAssertEqual(result3.items.count, 1)
   }
 
+  func testEditBounds() async throws {
+    try await SkipUnless.sourcekitdSupportsPlugin()
+    let sourcekitd = try await getSourceKitD()
+    let path = scratchFilePath()
+    _ = try await sourcekitd.openDocument(
+      path,
+      contents: "",
+      compilerArguments: [path]
+    )
+
+    let typeWithMethod = """
+      struct S {
+        static func foo() -> Int {}
+      }
+
+      """
+    var fullText = typeWithMethod
+
+    try await sourcekitd.editDocument(path, fromOffset: 0, length: 0, newContents: typeWithMethod)
+
+    let completion = """
+      S.
+      """
+    fullText += completion
+
+    try await sourcekitd.editDocument(path, fromOffset: typeWithMethod.utf8.count, length: 0, newContents: completion)
+
+    func testCompletion(file: StaticString = #filePath, line: UInt = #line) async throws {
+      let result = try await sourcekitd.completeOpen(
+        path: path,
+        position: Position(line: 3, utf16index: 2),
+        filter: "foo",
+        flags: []
+      )
+      XCTAssertGreaterThan(result.unfilteredResultCount, 1, file: file, line: line)
+      XCTAssertEqual(result.items.count, 1, file: file, line: line)
+    }
+    try await testCompletion()
+
+    // Bogus edits are ignored (negative offsets crash SourceKit itself so we don't test them here).
+    await assertThrowsError(
+      try await sourcekitd.editDocument(path, fromOffset: 0, length: 99999, newContents: "")
+    )
+    await assertThrowsError(
+      try await sourcekitd.editDocument(path, fromOffset: 99999, length: 1, newContents: "")
+    )
+    await assertThrowsError(
+      try await sourcekitd.editDocument(path, fromOffset: 99999, length: 0, newContents: "unrelated")
+    )
+    // SourceKit doesn't throw an error for a no-op edit.
+    try await sourcekitd.editDocument(path, fromOffset: 99999, length: 0, newContents: "")
+
+    try await sourcekitd.editDocument(path, fromOffset: 0, length: 0, newContents: "")
+    try await sourcekitd.editDocument(path, fromOffset: fullText.utf8.count, length: 0, newContents: "")
+
+    try await testCompletion()
+
+    let badCompletion = """
+      X.
+      """
+    fullText = fullText.dropLast(2) + badCompletion
+
+    try await sourcekitd.editDocument(path, fromOffset: fullText.utf8.count - 2, length: 2, newContents: badCompletion)
+
+    let result = try await sourcekitd.completeOpen(
+      path: path,
+      position: Position(line: 3, utf16index: 2),
+      filter: "foo",
+      flags: []
+    )
+    XCTAssertEqual(result.unfilteredResultCount, 0)
+    XCTAssertEqual(result.items.count, 0)
+  }
+
   func testDocumentation() async throws {
     try await SkipUnless.sourcekitdSupportsPlugin()
     try await SkipUnless.sourcekitdSupportsFullDocumentationInCompletion()