From da10fbbdcb32f8fae88082c424104786cedbba94 Mon Sep 17 00:00:00 2001
From: Alex Hoppen <ahoppen@apple.com>
Date: Wed, 30 Oct 2024 20:01:22 -0700
Subject: [PATCH] Fix a crash when trying to apply in edit that has out of line
 positions

`LineTable.replace` did not actually validate that the edit is in range. This could cause crashes of SourceKit-LSP if the editor is sending us bogus edits. Validate the position, like we do in all the other position conversions and log a fault if the edit is out-of-range.

While fixing this, I found a couple more places where line table accesses were not properly guarded. I migrated them to safe alternatives as well.

rdar://138962353
---
 Sources/SKSupport/LineTable.swift             | 66 ++++++++++--------
 .../Swift/CodeCompletionSession.swift         | 68 ++++++++++++-------
 Sources/SourceKitLSP/Swift/FoldingRange.swift |  9 +--
 Tests/SKSupportTests/LineTablePerfTests.swift |  6 +-
 Tests/SKSupportTests/SupportTests.swift       | 15 ++--
 Tests/SourceKitLSPTests/LifecycleTests.swift  | 27 ++++++++
 .../SwiftInterfaceTests.swift                 |  3 +-
 7 files changed, 127 insertions(+), 67 deletions(-)

diff --git a/Sources/SKSupport/LineTable.swift b/Sources/SKSupport/LineTable.swift
index a80294340..67296f993 100644
--- a/Sources/SKSupport/LineTable.swift
+++ b/Sources/SKSupport/LineTable.swift
@@ -36,25 +36,16 @@ package struct LineTable: Hashable, Sendable {
     }
   }
 
-  /// The number of lines.
-  @inlinable
-  package var count: Int { return impl.count }
-
-  /// Returns the given (zero-based) line as a Substring, including the newline.
-  ///
-  /// - parameter line: Line number (zero-based).
-  @inlinable
-  package subscript(line: Int) -> Substring {
-    return content[impl[line]..<(line == count - 1 ? content.endIndex : impl[line + 1])]
-  }
+  /// The number of lines in the line table.
+  package var lineCount: Int { impl.count }
 
   /// Translate String.Index to logical line/utf16 pair.
   package func lineAndUTF16ColumnOf(_ index: String.Index, fromLine: Int = 0) -> (line: Int, utf16Column: Int) {
-    precondition(0 <= fromLine && fromLine < count)
+    precondition(0 <= fromLine && fromLine < impl.count)
 
     // Binary search.
     var lower = fromLine
-    var upper = count
+    var upper = impl.count
     while true {
       let mid = lower + (upper - lower) / 2
       let lineStartIndex = impl[mid]
@@ -72,16 +63,6 @@ package struct LineTable: Hashable, Sendable {
   }
 }
 
-extension LineTable: RandomAccessCollection {
-  package var startIndex: Int {
-    return impl.startIndex
-  }
-
-  package var endIndex: Int {
-    return impl.endIndex
-  }
-}
-
 extension LineTable {
 
   // MARK: - Editing
@@ -101,8 +82,8 @@ extension LineTable {
     utf16Offset toOff: Int,
     with replacement: String
   ) {
-    let start = content.utf16.index(impl[fromLine], offsetBy: fromOff)
-    let end = content.utf16.index(impl[toLine], offsetBy: toOff)
+    let start = self.stringIndexOf(line: fromLine, utf16Column: fromOff)
+    let end = self.stringIndexOf(line: toLine, utf16Column: toOff)
 
     var newText = self.content
     newText.replaceSubrange(start..<end, with: replacement)
@@ -122,8 +103,14 @@ extension LineTable {
     utf16Length: Int,
     with replacement: String
   ) {
-    let start = content.utf16.index(impl[fromLine], offsetBy: fromOff)
-    let end = content.utf16.index(start, offsetBy: utf16Length)
+    let start = self.stringIndexOf(line: fromLine, utf16Column: fromOff)
+    let end: String.UTF16View.Index
+    if let endValue = content.utf16.index(start, offsetBy: utf16Length, limitedBy: content.endIndex) {
+      end = endValue
+    } else {
+      logger.fault("Range end is past end of file \(fromLine):\(fromOff) + \(utf16Length)")
+      end = content.endIndex
+    }
     let (toLine, toOff) = lineAndUTF16ColumnOf(end, fromLine: fromLine)
     self.replace(fromLine: fromLine, utf16Offset: fromOff, toLine: toLine, utf16Offset: toOff, with: replacement)
   }
@@ -158,7 +145,7 @@ extension LineTable {
       )
       return .beforeFirstLine
     }
-    guard line < count else {
+    guard line < impl.count else {
       logger.fault(
         """
         Line \(line) is out-of range (\(callerFile, privacy: .public):\(callerLine, privacy: .public))
@@ -166,7 +153,28 @@ extension LineTable {
       )
       return .afterLastLine
     }
-    return .line(self[line])
+    let start = impl[line]
+    let end: String.Index
+    if line + 1 < impl.count {
+      end = impl[line + 1]
+    } else {
+      end = content.endIndex
+    }
+
+    return .line(content[start..<end])
+  }
+
+  /// Extracts the contents of the line at the given index.
+  ///
+  /// If the line is out-of-bounds, returns `nil` and logs a fault.
+  @inlinable
+  package func line(at line: Int, callerFile: StaticString = #fileID, callerLine: UInt = #line) -> Substring? {
+    switch lineSlice(at: line, callerFile: callerFile, callerLine: callerLine) {
+    case .beforeFirstLine, .afterLastLine:
+      return nil
+    case .line(let line):
+      return line
+    }
   }
 
   /// Converts the given UTF-16-based `line:column`` position to a `String.Index`.
diff --git a/Sources/SourceKitLSP/Swift/CodeCompletionSession.swift b/Sources/SourceKitLSP/Swift/CodeCompletionSession.swift
index d419cfcfa..e14d8974f 100644
--- a/Sources/SourceKitLSP/Swift/CodeCompletionSession.swift
+++ b/Sources/SourceKitLSP/Swift/CodeCompletionSession.swift
@@ -417,46 +417,64 @@ class CodeCompletionSession {
     newText: String,
     snapshot: DocumentSnapshot
   ) -> TextEdit {
-    let textEditRangeStart: Position
+    let textEditRangeStart = computeCompletionTextEditStart(
+      completionPos: completionPos,
+      requestPosition: requestPosition,
+      utf8CodeUnitsToErase: utf8CodeUnitsToErase,
+      snapshot: snapshot
+    )
+    return TextEdit(range: textEditRangeStart..<requestPosition, newText: newText)
+  }
 
+  private func computeCompletionTextEditStart(
+    completionPos: Position,
+    requestPosition: Position,
+    utf8CodeUnitsToErase: Int,
+    snapshot: DocumentSnapshot
+  ) -> Position {
     // Compute the TextEdit
     if utf8CodeUnitsToErase == 0 {
       // Nothing to delete. Fast path and avoid UTF-8/UTF-16 conversions
-      textEditRangeStart = completionPos
+      return completionPos
     } else if utf8CodeUnitsToErase == 1 {
       // Fast path: Erasing a single UTF-8 byte code unit means we are also need to erase exactly one UTF-16 code unit, meaning we don't need to process the file contents
       if completionPos.utf16index >= 1 {
         // We can delete the character.
-        textEditRangeStart = Position(line: completionPos.line, utf16index: completionPos.utf16index - 1)
-      } else {
-        // Deleting the character would cross line boundaries. This is not supported by LSP.
-        // Fall back to ignoring utf8CodeUnitsToErase.
-        // If we discover that multi-lines replacements are often needed, we can add an LSP extension to support multi-line edits.
-        textEditRangeStart = completionPos
-      }
-    } else {
-      // We need to delete more than one text character. Do the UTF-8/UTF-16 dance.
-      assert(completionPos.line == requestPosition.line)
-      // Construct a string index for the edit range start by subtracting the UTF-8 code units to erase from the completion position.
-      let line = snapshot.lineTable[completionPos.line]
-      let deletionStartStringIndex = line.utf8.index(snapshot.index(of: completionPos), offsetBy: -utf8CodeUnitsToErase)
-
-      // Compute the UTF-16 offset of the deletion start range. If the start lies in a previous line, this will be negative
-      let deletionStartUtf16Offset = line.utf16.distance(from: line.startIndex, to: deletionStartStringIndex)
-
-      // Check if we are only deleting on one line. LSP does not support deleting over multiple lines.
-      if deletionStartUtf16Offset >= 0 {
-        // We are only deleting characters on the same line. Construct the corresponding text edit.
-        textEditRangeStart = Position(line: completionPos.line, utf16index: deletionStartUtf16Offset)
+        return Position(line: completionPos.line, utf16index: completionPos.utf16index - 1)
       } else {
         // Deleting the character would cross line boundaries. This is not supported by LSP.
         // Fall back to ignoring utf8CodeUnitsToErase.
         // If we discover that multi-lines replacements are often needed, we can add an LSP extension to support multi-line edits.
-        textEditRangeStart = completionPos
+        return completionPos
       }
     }
 
-    return TextEdit(range: textEditRangeStart..<requestPosition, newText: newText)
+    // We need to delete more than one text character. Do the UTF-8/UTF-16 dance.
+    assert(completionPos.line == requestPosition.line)
+    // Construct a string index for the edit range start by subtracting the UTF-8 code units to erase from the completion position.
+    guard let line = snapshot.lineTable.line(at: completionPos.line) else {
+      logger.fault("Code completion position is in out-of-range line \(completionPos.line)")
+      return completionPos
+    }
+    guard
+      let deletionStartStringIndex = line.utf8.index(
+        snapshot.index(of: completionPos),
+        offsetBy: -utf8CodeUnitsToErase,
+        limitedBy: line.utf8.startIndex
+      )
+    else {
+      // Deleting the character would cross line boundaries. This is not supported by LSP.
+      // Fall back to ignoring utf8CodeUnitsToErase.
+      // If we discover that multi-lines replacements are often needed, we can add an LSP extension to support multi-line edits.
+      logger.fault("UTF-8 code units to erase \(utf8CodeUnitsToErase) is before start of line")
+      return completionPos
+    }
+
+    // Compute the UTF-16 offset of the deletion start range.
+    let deletionStartUtf16Offset = line.utf16.distance(from: line.startIndex, to: deletionStartStringIndex)
+    precondition(deletionStartUtf16Offset >= 0)
+
+    return Position(line: completionPos.line, utf16index: deletionStartUtf16Offset)
   }
 }
 
diff --git a/Sources/SourceKitLSP/Swift/FoldingRange.swift b/Sources/SourceKitLSP/Swift/FoldingRange.swift
index e1107a0ee..5a1894a55 100644
--- a/Sources/SourceKitLSP/Swift/FoldingRange.swift
+++ b/Sources/SourceKitLSP/Swift/FoldingRange.swift
@@ -287,11 +287,12 @@ extension SwiftLanguageService {
 
 fileprivate extension LineTable {
   func isAtEndOfLine(_ position: Position) -> Bool {
-    guard position.line >= 0, position.line < self.count else {
+    guard let line = self.line(at: position.line) else {
       return false
     }
-    let line = self[position.line]
-    let suffixAfterPositionColumn = line[line.utf16.index(line.startIndex, offsetBy: position.utf16index)...]
-    return suffixAfterPositionColumn.allSatisfy(\.isNewline)
+    guard let index = line.utf16.index(line.startIndex, offsetBy: position.utf16index, limitedBy: line.endIndex) else {
+      return false
+    }
+    return line[index...].allSatisfy(\.isNewline)
   }
 }
diff --git a/Tests/SKSupportTests/LineTablePerfTests.swift b/Tests/SKSupportTests/LineTablePerfTests.swift
index 72a19a794..b5840249f 100644
--- a/Tests/SKSupportTests/LineTablePerfTests.swift
+++ b/Tests/SKSupportTests/LineTablePerfTests.swift
@@ -87,9 +87,9 @@ final class LineTablePerfTests: PerfTestCase {
       self.startMeasuring()
 
       for _ in 1...iterations {
-        let line = (0..<(t.count - 1)).randomElement(using: &lcg) ?? 0
-        let col = (0..<t[line].utf16.count).randomElement(using: &lcg) ?? 0
-        let len = t[line].isEmpty ? 0 : Bool.random() ? 1 : 0
+        let line = (0..<(t.lineCount - 1)).randomElement(using: &lcg) ?? 0
+        let col = (0..<t.line(at: line)!.utf16.count).randomElement(using: &lcg) ?? 0
+        let len = t.line(at: line)!.isEmpty ? 0 : Bool.random() ? 1 : 0
         var newText = String(characters.randomElement(using: &lcg)!)
         if len == 1 && Bool.random(using: &lcg) {
           newText = ""  // deletion
diff --git a/Tests/SKSupportTests/SupportTests.swift b/Tests/SKSupportTests/SupportTests.swift
index 397d4a931..38811a53b 100644
--- a/Tests/SKSupportTests/SupportTests.swift
+++ b/Tests/SKSupportTests/SupportTests.swift
@@ -14,17 +14,22 @@ import SKSupport
 import TSCBasic
 import XCTest
 
-final class SupportTests: XCTestCase {
+fileprivate extension LineTable {
+  var lines: [Substring] {
+    (0..<lineCount).map { line(at: $0)! }
+  }
+}
 
+final class SupportTests: XCTestCase {
   func checkLines(_ string: String, _ expected: [String], file: StaticString = #filePath, line: UInt = #line) {
     let table = LineTable(string)
-    XCTAssertEqual(table.map { String($0) }, expected, file: file, line: line)
+    XCTAssertEqual(table.lines.map(String.init), expected, file: file, line: line)
   }
 
   func checkOffsets(_ string: String, _ expected: [Int], file: StaticString = #filePath, line: UInt = #line) {
     let table = LineTable(string)
     XCTAssertEqual(
-      table.map { string.utf8.distance(from: string.startIndex, to: $0.startIndex) },
+      table.lines.map { string.utf8.distance(from: string.startIndex, to: $0.startIndex) },
       expected,
       file: file,
       line: line
@@ -69,8 +74,8 @@ final class SupportTests: XCTestCase {
     checkOffsets("\n\u{10000}b", [0, 1])
     checkOffsets("\n\u{10000}b\nc", [0, 1, 7])
 
-    XCTAssertEqual(LineTable("")[0], "")
-    XCTAssertEqual(LineTable("\n")[1], "")
+    XCTAssertEqual(LineTable("").line(at: 0), "")
+    XCTAssertEqual(LineTable("\n").line(at: 1), "")
   }
 
   func checkLineAndColumns(
diff --git a/Tests/SourceKitLSPTests/LifecycleTests.swift b/Tests/SourceKitLSPTests/LifecycleTests.swift
index d7ea3f6ad..79ac94318 100644
--- a/Tests/SourceKitLSPTests/LifecycleTests.swift
+++ b/Tests/SourceKitLSPTests/LifecycleTests.swift
@@ -135,4 +135,31 @@ final class LifecycleTests: XCTestCase {
     )
     XCTAssertGreaterThan(symbolInfo.count, 0)
   }
+
+  func testEditWithOutOfRangeLine() async throws {
+    let testClient = try await TestSourceKitLSPClient()
+    let uri = DocumentURI(for: .swift)
+    testClient.openDocument("", uri: uri)
+
+    // Check that we don't crash.
+    testClient.send(
+      DidChangeTextDocumentNotification(
+        textDocument: VersionedTextDocumentIdentifier(uri, version: 2),
+        contentChanges: [TextDocumentContentChangeEvent(range: Range(Position(line: 2, utf16index: 0)), text: "new")]
+      )
+    )
+  }
+
+  func testEditWithOutOfRangeColumn() async throws {
+    let testClient = try await TestSourceKitLSPClient()
+    let uri = DocumentURI(for: .swift)
+    testClient.openDocument("", uri: uri)
+
+    testClient.send(
+      DidChangeTextDocumentNotification(
+        textDocument: VersionedTextDocumentIdentifier(uri, version: 2),
+        contentChanges: [TextDocumentContentChangeEvent(range: Range(Position(line: 0, utf16index: 4)), text: "new")]
+      )
+    )
+  }
 }
diff --git a/Tests/SourceKitLSPTests/SwiftInterfaceTests.swift b/Tests/SourceKitLSPTests/SwiftInterfaceTests.swift
index bf1ee04bf..02dec2a37 100644
--- a/Tests/SourceKitLSPTests/SwiftInterfaceTests.swift
+++ b/Tests/SourceKitLSPTests/SwiftInterfaceTests.swift
@@ -200,6 +200,7 @@ private func assertSystemSwiftInterface(
   // load contents of swiftinterface
   let contents = try XCTUnwrap(location.uri.fileURL.flatMap({ try String(contentsOf: $0, encoding: .utf8) }))
   let lineTable = LineTable(contents)
-  let destinationLine = lineTable[location.range.lowerBound.line].trimmingCharacters(in: .whitespaces)
+  let destinationLine = try XCTUnwrap(lineTable.line(at: location.range.lowerBound.line))
+    .trimmingCharacters(in: .whitespaces)
   XCTAssert(destinationLine.hasPrefix(linePrefix), "Full line was: '\(destinationLine)'", line: line)
 }