[clangd] Sanity-check array sizes read from disk before allocating them.

[sammccall]

Summary

Previously a corrupted index shard could cause us to resize arrays to an
arbitrary int32. This tends to be a huge number, and can render the
system unresponsive.

Instead, cap this at the amount of data that might reasonably be read
(e.g. the #bytes in the file). If the specified length is more than that,
assume the data is corrupt.

[kadircet]

clang-tools-extra/clangd/index/Serialization.cpp

121

regarding minsizes, i suppose the idea was to pass ElementSizeInBytes for containers ? I am OK with the overshooting if you don't want to make this more detailed, but can we drop the param?

211

in theory, this could also trip us over. zlib::uncompress does a reserve using this value. it is harder to come up with an upper bound here, but https://zlib.net/zlib_tech.html#:~:text=Maximum%20Compression%20Factor&text=(The%20test%20case%20was%20a,sources)%20is%201032%3A1. says the theoretical limit is 1032:1.
Maybe enforce that?

Also to make this more similar to others, it looks like zlib::uncompress also has an overload taking a const char *.

521

unrelated to this patch, but it feels like we are checking for the error at the wrong place ?

clang-tools-extra/clangd/unittests/SerializationTests.cpp

320

should we rather ADD_FAILURE ?

345

why not ASSERT_TRUE(readIndexFile(..)) ? (or llvm::cantFail)

364

can we use In.Sources.begin()->Digest instead?

[sammccall]

clang-tools-extra/clangd/index/Serialization.cpp

121

Yeah fair enough - I thought having the param made the logic easier to follow but actually passing it made the code too fragile.

Just inlined badSize() and hardcoded MinSize=1, we can always extract it again later.

521

agree. Will fix in a different commit

clang-tools-extra/clangd/unittests/SerializationTests.cpp

320

No, because rlimit can legitimately fail (e.g. we're already using more ram than the limit we tried to set).

If rlimit fails then we may not be able to perform the test meaningfully, so we could skip it in that case (i.e. decide to *pass* immediately).

345

Whoops, this was leftover experimentation code.

364

Hmm, we could, but it has the wrong data type (uint8_t vs char) so we need a reinterpret_cast... not sure it's clearer.

[sammccall]

Address comments

[kadircet]

thanks, lgtm!

[eugenis]

Hi Sam,

this patch is failing on the ubsan bot with:

[ RUN      ] SerializationTest.NoCrashOnBadArraySize
/b/sanitizer-x86_64-linux-fast/build/llvm-project/clang-tools-extra/clangd/index/Serialization.cpp:90:26: runtime error: left shift of 127 by 28 places cannot be represented in type 'int'
    #0 0x392dd57 in clang::clangd::(anonymous namespace)::Reader::consumeVar() /b/sanitizer-x86_64-linux-fast/build/llvm-project/clang-tools-extra/clangd/index/Serialization.cpp:90:26
    #1 0x392dc68 in bool clang::clangd::(anonymous namespace)::Reader::consumeSize<std::__1::vector<llvm::StringRef, std::__1::allocator<llvm::StringRef> > >(std::__1::vector<llvm::StringRef, std::__1::allocator<llvm::StringRef> >&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/clang-tools-extra/clangd/index/Serialization.cpp:113:17
    #2 0x3926d6a in readIncludeGraphNode /b/sanitizer-x86_64-linux-fast/build/llvm-project/clang-tools-extra/clangd/index/Serialization.cpp:275:13
    #3 0x3926d6a in clang::clangd::(anonymous namespace)::readRIFF(llvm::StringRef) /b/sanitizer-x86_64-linux-fast/build/llvm-project/clang-tools-extra/clangd/index/Serialization.cpp:474:18
    #4 0x3926111 in clang::clangd::readIndexFile(llvm::StringRef) /b/sanitizer-x86_64-linux-fast/build/llvm-project/clang-tools-extra/clangd/index/Serialization.cpp:676:12
    #5 0x1e5b77c in clang::clangd::(anonymous namespace)::SerializationTest_NoCrashOnBadArraySize_Test::TestBody() /b/sanitizer-x86_64-linux-fast/build/llvm-project/clang-tools-extra/clangd/unittests/SerializationTests.cpp:380:24
    #6 0x20589f9 in testing::Test::Run() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/src/gtest.cc:2474:5
    #7 0x205993b in testing::TestInfo::Run() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/src/gtest.cc:2656:11
    #8 0x205a4e2 in testing::TestCase::Run() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/src/gtest.cc:2774:28
    #9 0x20619b2 in testing::internal::UnitTestImpl::RunAllTests() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/src/gtest.cc:4649:43
    #10 0x20613c9 in testing::UnitTest::Run() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/src/gtest.cc:4257:10
    #11 0x2051273 in RUN_ALL_TESTS /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/googletest/include/gtest/gtest.h:2233:46
    #12 0x2051273 in main /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/utils/unittest/UnitTestMain/TestMain.cpp:50:10
    #13 0x7fd11479d09a in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x2409a)
    #14 0x1a330d9 in _start (/b/sanitizer-x86_64-linux-fast/build/llvm_build_ubsan/tools/clang/tools/extra/clangd/unittests/ClangdTests+0x1a330d9)

It looks like the corrupt input in your test case is triggering a preexisting bug in clangd.

[kadircet]

oops, this is actually a problem of the test. it is providing an invalid encoding :/ sent out https://reviews.llvm.org/D91405 to fix the issue.

[sammccall]

In D91258#2392646, @eugenis wrote:

Hi Sam,

this patch is failing on the ubsan bot with:

[ RUN      ] SerializationTest.NoCrashOnBadArraySize
/b/sanitizer-x86_64-linux-fast/build/llvm-project/clang-tools-extra/clangd/index/Serialization.cpp:90:26: runtime error: left shift of 127 by 28 places cannot be represented in type 'int'
    #0 0x392dd57 in clang::clangd::(anonymous namespace)::Reader::consumeVar() /b/sanitizer-x86_64-linux-fast/build/llvm-project/clang-tools-extra/clangd/index/Serialization.cpp:90:26

Thanks Evgenii! (It's really surprising to me this is using signed arithmetic!)

Is the sanitizer strictly correct here that this is UB? From my reading only C requires the result of signed shifts to be representable, not C++. https://eel.is/c++draft/expr.shift
So this should be well-defined with the same result as the unsigned arithmetic (assuming 2s complement, I guess).

(We'll fix this to use unsigned arithmetic in any case)

clang-tools-extra/clangd/index/Serialization.cpp

 #include "RIFF.h"
 #include "SymbolLocation.h"
 #include "SymbolOrigin.h"
 #include "dex/Dex.h"
 #include "support/Logger.h"
 #include "support/Trace.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Compiler.h"
 #include "llvm/Support/Compression.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/raw_ostream.h"
 #include <vector>
 
 namespace clang {
 namespace clangd {
llvm::StringRef consumeString(llvm::ArrayRef<llvm::StringRef> Strings) {
     }
     return Strings[StringIndex];
   }
 
   SymbolID consumeID() {
     llvm::StringRef Raw = consume(SymbolID::RawSize); // short if truncated.
     return LLVM_UNLIKELY(err()) ? SymbolID() : SymbolID::fromRaw(Raw);
   }
+
+  // Read a varint (as consumeVar) and resize the container accordingly.
+  // If the size is invalid, return false and mark an error.
+  // (The caller should abort in this case).
+  template <typename T> LLVM_NODISCARD bool consumeSize(T &Container) {
+    auto Size = consumeVar();
+    // Conservatively assume each element is at least one byte.
+    if (Size > (End - Begin)) {
+      Err = true;
+      return false;
+    }
+    Container.resize(Size);
+    return true;
+  }
 };
 
 void write32(uint32_t I, llvm::raw_ostream &OS) {
   char Buf[4];
   llvm::support::endian::write32le(Buf, I);
   OS.write(Buf, sizeof(Buf));
 }
 

 
 struct StringTableIn {
   llvm::BumpPtrAllocator Arena;
   std::vector<llvm::StringRef> Strings;
 };
 
 llvm::Expected<StringTableIn> readStringTable(llvm::StringRef Data) {
   Reader R(Data);
   size_t UncompressedSize = R.consume32();
   if (R.err())
     return error("Truncated string table");
 
   llvm::StringRef Uncompressed;
   llvm::SmallString<1> UncompressedStorage;
   if (UncompressedSize == 0) // No compression
     Uncompressed = R.rest();
   else if (llvm::zlib::isAvailable()) {

 
 IncludeGraphNode readIncludeGraphNode(Reader &Data,
                                       llvm::ArrayRef<llvm::StringRef> Strings) {
   IncludeGraphNode IGN;
   IGN.Flags = static_cast<IncludeGraphNode::SourceFlag>(Data.consume8());
   IGN.URI = Data.consumeString(Strings);
   llvm::StringRef Digest = Data.consume(IGN.Digest.size());
   std::copy(Digest.bytes_begin(), Digest.bytes_end(), IGN.Digest.begin());
-  IGN.DirectIncludes.resize(Data.consumeVar());
+  if (!Data.consumeSize(IGN.DirectIncludes))
+    return IGN;
   for (llvm::StringRef &Include : IGN.DirectIncludes)
     Include = Data.consumeString(Strings);
   return IGN;
 }
 
 void writeIncludeGraphNode(const IncludeGraphNode &IGN,
                            const StringTableOut &Strings,
                            llvm::raw_ostream &OS) {
Symbol readSymbol(Reader &Data, llvm::ArrayRef<llvm::StringRef> Strings) {
   Sym.References = Data.consumeVar();
   Sym.Flags = static_cast<Symbol::SymbolFlag>(Data.consume8());
   Sym.Origin = static_cast<SymbolOrigin>(Data.consume8());
   Sym.Signature = Data.consumeString(Strings);
   Sym.CompletionSnippetSuffix = Data.consumeString(Strings);
   Sym.Documentation = Data.consumeString(Strings);
   Sym.ReturnType = Data.consumeString(Strings);
   Sym.Type = Data.consumeString(Strings);
-  Sym.IncludeHeaders.resize(Data.consumeVar());
+  if (!Data.consumeSize(Sym.IncludeHeaders))
+    return Sym;
   for (auto &I : Sym.IncludeHeaders) {
     I.IncludeHeader = Data.consumeString(Strings);
     I.References = Data.consumeVar();
   }
   return Sym;
 }
 
 // REFS ENCODING
for (const auto &Ref : Refs) {
     OS << Ref.Container.raw();
   }
 }
 
 std::pair<SymbolID, std::vector<Ref>>
 readRefs(Reader &Data, llvm::ArrayRef<llvm::StringRef> Strings) {
   std::pair<SymbolID, std::vector<Ref>> Result;
   Result.first = Data.consumeID();
-  Result.second.resize(Data.consumeVar());
+  if (!Data.consumeSize(Result.second))
+    return Result;
   for (auto &Ref : Result.second) {
     Ref.Kind = static_cast<RefKind>(Data.consume8());
     Ref.Location = readLocation(Data, Strings);
     Ref.Container = Data.consumeID();
   }
   return Result;
 }
 
void writeCompileCommand(const InternedCompileCommand &Cmd,
   for (llvm::StringRef C : Cmd.CommandLine)
     writeVar(Strings.index(C), CmdOS);
 }
 
 InternedCompileCommand
 readCompileCommand(Reader CmdReader, llvm::ArrayRef<llvm::StringRef> Strings) {
   InternedCompileCommand Cmd;
   Cmd.Directory = CmdReader.consumeString(Strings);
-  Cmd.CommandLine.resize(CmdReader.consumeVar());
+  if (!CmdReader.consumeSize(Cmd.CommandLine))
+    return Cmd;
   for (llvm::StringRef &C : Cmd.CommandLine)
     C = CmdReader.consumeString(Strings);
   return Cmd;
 }
 
 // FILE ENCODING
 // A file is a RIFF chunk with type 'CdIx'.
 // It contains the sections:
while (!RelationsReader.eof()) {
       Relations.insert(Relation);
     }
     if (RelationsReader.err())
       return error("malformed or truncated relations");
     Result.Relations = std::move(Relations).build();
   }
   if (Chunks.count("cmdl")) {
     Reader CmdReader(Chunks.lookup("cmdl"));
     if (CmdReader.err())
       return error("malformed or truncated commandline section");
     InternedCompileCommand Cmd =
         readCompileCommand(CmdReader, Strings->Strings);
     Result.Cmd.emplace();
     Result.Cmd->Directory = std::string(Cmd.Directory);
     Result.Cmd->CommandLine.reserve(Cmd.CommandLine.size());
     for (llvm::StringRef C : Cmd.CommandLine)
       Result.Cmd->CommandLine.emplace_back(C);

clang-tools-extra/clangd/unittests/SerializationTests.cpp

 //===-- SerializationTests.cpp - Binary and YAML serialization unit tests -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Headers.h"
+#include "RIFF.h"
 #include "index/Index.h"
 #include "index/Serialization.h"
+#include "support/Logger.h"
 #include "clang/Tooling/CompilationDatabase.h"
+#include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#ifdef LLVM_ON_UNIX
+#include <sys/resource.h>
+#endif
 
-using ::testing::_;
-using ::testing::AllOf;
 using ::testing::ElementsAre;
 using ::testing::Pair;
 using ::testing::UnorderedElementsAre;
 using ::testing::UnorderedElementsAreArray;
 
 namespace clang {
 namespace clangd {
 namespace {
Out.Cmd = &Cmd;
     const tooling::CompileCommand &SerializedCmd = In->Cmd.getValue();
     EXPECT_EQ(SerializedCmd.CommandLine, Cmd.CommandLine);
     EXPECT_EQ(SerializedCmd.Directory, Cmd.Directory);
     EXPECT_NE(SerializedCmd.Filename, Cmd.Filename);
     EXPECT_NE(SerializedCmd.Heuristic, Cmd.Heuristic);
     EXPECT_NE(SerializedCmd.Output, Cmd.Output);
   }
 }
+
+#if LLVM_ON_UNIX // rlimit is part of POSIX
+class ScopedMemoryLimit {
+  struct rlimit OriginalLimit;
+  bool Succeeded = false;
+
+public:
+  ScopedMemoryLimit(rlim_t Bytes) {
+    if (!getrlimit(RLIMIT_AS, &OriginalLimit)) {
+      struct rlimit NewLimit = OriginalLimit;
+      NewLimit.rlim_cur = Bytes;
+      Succeeded = !setrlimit(RLIMIT_AS, &NewLimit);
+    }
+    if (!Succeeded)
+      log("Failed to set rlimit");
+  }
+
+  ~ScopedMemoryLimit() {
+    if (Succeeded)
+      setrlimit(RLIMIT_AS, &OriginalLimit);
+  }
+};
+#else
+class ScopedMemoryLimit {
+public:
+  ScopedMemoryLimit(unsigned Bytes) { log("rlimit unsupported"); }
+};
+#endif
+
+// Test that our deserialization detects invalid array sizes without allocating.
+// If this detection fails, the test should allocate a huge array and crash.
+TEST(SerializationTest, NoCrashOnBadArraySize) {
+  // This test is tricky because we need to construct a subtly invalid file.
+  // First, create a valid serialized file.
+  auto In = readIndexFile(YAML);
+  ASSERT_FALSE(!In) << In.takeError();
+  IndexFileOut Out(*In);
+  Out.Format = IndexFileFormat::RIFF;
+  std::string Serialized = llvm::to_string(Out);
+
+  // Low-level parse it again and find the `srcs` chunk we're going to corrupt.
+  auto Parsed = riff::readFile(Serialized);
+  ASSERT_FALSE(!Parsed) << Parsed.takeError();
+  auto Srcs = llvm::find_if(Parsed->Chunks, [](riff::Chunk C) {
+    return C.ID == riff::fourCC("srcs");
+  });
+  ASSERT_NE(Srcs, Parsed->Chunks.end());
+
+  // Srcs consists of a sequence of IncludeGraphNodes. In our case, just one.
+  // The node has:
+  //  - 1 byte: flags (1)
+  //  - varint(stringID): URI
+  //  - 8 byte: file digest
+  //  - varint: DirectIncludes.length
+  //  - repeated varint(stringID): DirectIncludes
+  // We want to set DirectIncludes.length to a huge number.
+  // The offset isn't trivial to find, so we use the file digest.
+  std::string FileDigest = llvm::fromHex("EED8F5EAF25C453C");
+  unsigned Pos = Srcs->Data.find_first_of(FileDigest);
+  ASSERT_NE(Pos, StringRef::npos) << "Couldn't locate file digest";
+  Pos += FileDigest.size();
+
+  // Varints are little-endian base-128 numbers, where the top-bit of each byte
+  // indicates whether there are more. 8fffffff7f -> 0xffffffff.
+  std::string CorruptSrcs =
+      (Srcs->Data.take_front(Pos) + llvm::fromHex("8fffffff7f") +
+       "some_random_garbage")
+          .str();
+  Srcs->Data = CorruptSrcs;
+
+  // Try to crash rather than hang on large allocation.
+  ScopedMemoryLimit MemLimit(1000 * 1024 * 1024); // 1GB
+
+  std::string CorruptFile = llvm::to_string(*Parsed);
+  auto CorruptParsed = readIndexFile(CorruptFile);
+  ASSERT_TRUE(!CorruptParsed);
+  EXPECT_EQ(llvm::toString(CorruptParsed.takeError()),
+            "malformed or truncated include uri");
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang