From a2f7352f3e809cb1b267e769d00ea84e4ef46bf0 Mon Sep 17 00:00:00 2001
From: Kadir Cetinkaya <kadircet@google.com>
Date: Wed, 24 May 2023 14:22:43 +0200
Subject: [PATCH] [clangd] Dont run raw-lexer for OOB source locations

We can get stale source locations from preamble, make sure we don't
access those locations without checking first.

Fixes https://github.com/clangd/clangd/issues/1636.

Differential Revision: https://reviews.llvm.org/D151321
---
 clang-tools-extra/clangd/Diagnostics.cpp           | 17 +++++++++---
 .../clangd/unittests/PreambleTests.cpp             | 31 ++++++++++++++++++++++
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/clang-tools-extra/clangd/Diagnostics.cpp b/clang-tools-extra/clangd/Diagnostics.cpp
index b708f9c..4c5def3 100644
--- a/clang-tools-extra/clangd/Diagnostics.cpp
+++ b/clang-tools-extra/clangd/Diagnostics.cpp
@@ -96,7 +96,8 @@ bool locationInRange(SourceLocation L, CharSourceRange R,
 
 // Clang diags have a location (shown as ^) and 0 or more ranges (~~~~).
 // LSP needs a single range.
-Range diagnosticRange(const clang::Diagnostic &D, const LangOptions &L) {
+std::optional<Range> diagnosticRange(const clang::Diagnostic &D,
+                                     const LangOptions &L) {
   auto &M = D.getSourceManager();
   auto PatchedRange = [&M](CharSourceRange &R) {
     R.setBegin(translatePreamblePatchLocation(R.getBegin(), M));
@@ -115,13 +116,18 @@ Range diagnosticRange(const clang::Diagnostic &D, const LangOptions &L) {
     if (locationInRange(Loc, R, M))
       return halfOpenToRange(M, PatchedRange(R));
   }
+  // Source locations from stale preambles might become OOB.
+  // FIXME: These diagnostics might point to wrong locations even when they're
+  // not OOB.
+  auto [FID, Offset] = M.getDecomposedLoc(Loc);
+  if (Offset > M.getBufferData(FID).size())
+    return std::nullopt;
   // If the token at the location is not a comment, we use the token.
   // If we can't get the token at the location, fall back to using the location
   auto R = CharSourceRange::getCharRange(Loc);
   Token Tok;
-  if (!Lexer::getRawToken(Loc, Tok, M, L, true) && Tok.isNot(tok::comment)) {
+  if (!Lexer::getRawToken(Loc, Tok, M, L, true) && Tok.isNot(tok::comment))
     R = CharSourceRange::getTokenRange(Tok.getLocation(), Tok.getEndLoc());
-  }
   return halfOpenToRange(M, PatchedRange(R));
 }
 
@@ -697,7 +703,10 @@ void StoreDiags::HandleDiagnostic(DiagnosticsEngine::Level DiagLevel,
     SourceLocation PatchLoc =
         translatePreamblePatchLocation(Info.getLocation(), SM);
     D.InsideMainFile = isInsideMainFile(PatchLoc, SM);
-    D.Range = diagnosticRange(Info, *LangOpts);
+    if (auto DRange = diagnosticRange(Info, *LangOpts))
+      D.Range = *DRange;
+    else
+      D.Severity = DiagnosticsEngine::Ignored;
     auto FID = SM.getFileID(Info.getLocation());
     if (const auto FE = SM.getFileEntryRefForID(FID)) {
       D.File = FE->getName().str();
diff --git a/clang-tools-extra/clangd/unittests/PreambleTests.cpp b/clang-tools-extra/clangd/unittests/PreambleTests.cpp
index 23ba0fc..4f2cc3e 100644
--- a/clang-tools-extra/clangd/unittests/PreambleTests.cpp
+++ b/clang-tools-extra/clangd/unittests/PreambleTests.cpp
@@ -829,6 +829,37 @@ x>)");
     auto AST = createPatchedAST(Code.code(), NewCode.code());
     EXPECT_THAT(*AST->getDiagnostics(), IsEmpty());
   }
+  {
+    Annotations Code(R"(
+#ifndef FOO
+#define FOO
+void foo();
+#endif)");
+    // This code will emit a diagnostic for unterminated #ifndef (as stale
+    // preamble has the conditional but main file doesn't terminate it).
+    // We shouldn't emit any diagnotiscs (and shouldn't crash).
+    Annotations NewCode("");
+    auto AST = createPatchedAST(Code.code(), NewCode.code());
+    EXPECT_THAT(*AST->getDiagnostics(), IsEmpty());
+  }
+  {
+    Annotations Code(R"(
+#ifndef FOO
+#define FOO
+void foo();
+#endif)");
+    // This code will emit a diagnostic for unterminated #ifndef (as stale
+    // preamble has the conditional but main file doesn't terminate it).
+    // We shouldn't emit any diagnotiscs (and shouldn't crash).
+    // FIXME: Patch/ignore diagnostics in such cases.
+    Annotations NewCode(R"(
+i[[nt]] xyz;
+    )");
+    auto AST = createPatchedAST(Code.code(), NewCode.code());
+    EXPECT_THAT(
+        *AST->getDiagnostics(),
+        ElementsAre(Diag(NewCode.range(), "pp_unterminated_conditional")));
+  }
 }
 
 MATCHER_P2(Mark, Range, Text, "") {
-- 
2.7.4